tmp

Author: alanwaketan

test/test_operations.py

@@ -1867,6 +1867,16 @@ def test_patched_linear_1D_bias(self):
     self.assertTrue(
         torch.allclose(linear.bias.grad.cpu(), linear_cpu.bias.grad))
 
+  def test_tpu_custom_call(self):
+    payload = "{\"custom_call_config\": {\"body\": \"TUzvUgFNTElSMTguMC4wZ2l0AAExCQEDBQcBAwkDBQMLBQkNDxETBwkVFxkbA31bEwFXBw8LCw8PEwsLGwszCw8LMwsLCwsLCwsTHwsTHxsPCxNlDwsTDxMPIwtTEwUFYVkBExsfDw8fGwcLGwJeBB8dMw0BAQUdHT0/EQcBFw8VAQUfBSEDBRUXGRsFIyMPEREBAAAAAAAAAAUlEQUFBScDCyEFIyUnCykLKy0FKQUrDREFLQUvBTEFMwMDBzElAQkIAAAABTUDAwc3JQEJAAAAAAMFEQUHOxENAQU3Fw8TAWFmZmluZV9tYXA8KGQwKSAtPiAoZDApPgAdRQ0FOQMDHUkRBxUDAx1NEQcJAwdRUxUXGRsFOyMFCSEBAAAAAQAAAAEAAAABAAAAAwMRBSN0cHUubWVtb3J5X3NwYWNlPHZtZW0+ACN0cHUudGlsZWQ8KDEyOCksWzFdPgAnBSECBAUXVwMCBAVZAQICAQIEF1cDAgQFQScFIQIEDwMBCQUHAwMDAQQSAgUBEAEHAwEFAxEBHwcDIT8HAwEDAQMBDwMDLwMBDwMDNQMBDwMJOQMNBwYBAwkDAQcGAQMJAwMHBgEDCQMFCQcJEwMBBQ0LCQcJEwMBBQ8LEwZDAwEFExULAgMDAREHA0cDCwUZCREHA0sDCwUZBxUGAwMLBRsdDQUDTwkXEQsfBQEBVQYDAQUBAOoFPSkLOTsLEyMhHSkVHxsVQw0LCwsTDQsLKQ8PDQkLEWJ1aWx0aW4AZnVuYwB0cHUAYXJpdGgAbW9kdWxlAHJldHVybgBlcmFzZV9tZW1yZWZfbGF5b3V0AGxvYWQAaW90YQBzdG9yZQBjb25zdGFudABjbXBpAGFkZGkAYW5kaQB2YWx1ZQAvaG9tZS9qd3Rhbi9wYWxsYXMvcGFsbGFzX2FkZC5weQBpbl9sYXlvdXQAc3VibGFuZV9tYXNrAHN1YmxhbmVfc3RyaWRlAHByZWRpY2F0ZQBkaW1lbnNpb25fc2VtYW50aWNzAGZ1bmN0aW9uX3R5cGUAc2NhbGFyX3ByZWZldGNoAHNjcmF0Y2hfb3BlcmFuZHMAc3ltX25hbWUAbWFpbgAvc3dhcFtpbmRleGVkX2RpbXM9KEZhbHNlLCldAC9nZXRbaW5kZXhlZF9kaW1zPShGYWxzZSwpXQAvYWRkAG9wZXJhbmRTZWdtZW50U2l6ZXMA\"}}"
+
+    x = torch.arange(8, dtype=torch.int).to("xla")
+    y = torch.arange(8, dtype=torch.int).to("xla")
+    output = torch.arange(8, dtype=torch.int).to("xla")
+
+    torch_xla._XLAC._xla_tpu_custom_call_(output, x, y, payload)
+    print(output)
+
 
 class MNISTComparator(nn.Module):
 

torch_xla/csrc/init_python_bindings.cpp

@@ -2170,6 +2170,11 @@ void InitXlaModuleBindings(py::module m) {
         [](at::Tensor& self, const at::Tensor& source) -> at::Tensor& {
           return XLANativeFunctions::set_(self, source);
         });
+  m.def("_xla_tpu_custom_call_",
+      [](at::Tensor& output, const at::Tensor& x, const at::Tensor& y, const std::string& payload) {
+        auto x_output = bridge::GetXlaTensor(output);
+        return tensor_methods::tpu_custom_call_(x_output, bridge::GetXlaTensor(x), bridge::GetXlaTensor(y), payload);
+      });
   m.def("_set_xla_custom_op_name_prefix",
         [](const at::Tensor& input, const std::string& op_name_prefix,
            size_t max_call_stack_depth) -> bool {

torch_xla/csrc/ops/xla_ops.cpp

@@ -33,5 +33,6 @@ const OpKindWrapper xla_tensor_data("xla::tensor_data");
 const OpKindWrapper xla_unselect("xla::unselect");
 const OpKindWrapper xla_update_slice("xla::update_slice");
 const OpKindWrapper xla_custom_sharding("xla::custom_sharding");
+const OpKindWrapper xla_tpu_custom_call("xla::tpu_custom_call");
 
 }  // namespace torch_xla

torch_xla/csrc/ops/xla_ops.h

@@ -58,6 +58,7 @@ extern const OpKindWrapper xla_tensor_data;
 extern const OpKindWrapper xla_unselect;
 extern const OpKindWrapper xla_update_slice;
 extern const OpKindWrapper xla_custom_sharding;
+extern const OpKindWrapper xla_tpu_custom_call;
 
 }  // namespace torch_xla
 

torch_xla/csrc/tensor_methods.cpp

@@ -41,6 +41,7 @@
 #include "torch_xla/csrc/ops/cumprod.h"
 #include "torch_xla/csrc/ops/cumsum.h"
 #include "torch_xla/csrc/ops/custom_sharding.h"
+#include "torch_xla/csrc/ops/tpu_custom_call.h"
 #include "torch_xla/csrc/ops/dequant_tensor.h"
 #include "torch_xla/csrc/ops/device_data.h"
 #include "torch_xla/csrc/ops/diagonal.h"
@@ -523,6 +524,11 @@ void custom_sharding_(
   input->SetShardingSpec(*sharding_spec);
 }
 
+void tpu_custom_call_(XLATensorPtr& output, const XLATensorPtr& x, const XLATensorPtr& y, const std::string& payload) {
+  output->SetInPlaceIrValue(
+      torch::lazy::MakeNode<TpuCustomCall>(x->GetIrValue(), y->GetIrValue(), payload));
+}
+
 XLATensorPtr get_dimensions_size(const XLATensorPtr& input,
                                  std::vector<int64_t> dimensions) {
   return input->CreateFrom(torch::lazy::MakeNode<GetDimensionsSize>(

torch_xla/csrc/tensor_methods.h

@@ -82,6 +82,8 @@ std::pair<XLATensorPtr, torch::lazy::Value> collective_permute(
 void custom_sharding_(const XLATensorPtr& input,
                       const std::shared_ptr<XLATensor::ShardingSpec>& spec);
 
+void tpu_custom_call_(XLATensorPtr& output, const XLATensorPtr& x, const XLATensorPtr& y, const std::string& payload);
+
 XLATensorPtr get_dimensions_size(const XLATensorPtr& input,
                                  std::vector<int64_t> dimensions);
 

torch_xla/csrc/xla_lower_util.cpp

@@ -1224,4 +1224,10 @@ xla::XlaOp BuildCustomSharding(const xla::XlaOp& input) {
                          {input}, ShapeHelper::ShapeOfXlaOp(input));
 }
 
+xla::XlaOp BuildTpuCustomCall(const xla::XlaOp& x, const xla::XlaOp& y, const std::string& payload) {
+  return xla::CustomCallWithLayout(x.builder(), /*call_target_name=*/"tpu_custom_call",
+                         {x, y}, ShapeHelper::ShapeOfXlaOp(x) /*TODO: update later*/,
+                         {ShapeHelper::ShapeOfXlaOp(x), ShapeHelper::ShapeOfXlaOp(y)}, payload);
+}
+
 }  // namespace torch_xla

torch_xla/csrc/xla_lower_util.h

@@ -150,6 +150,8 @@ xla::XlaOp BuildCdistForward(xla::XlaOp x1, xla::XlaOp x2, xla::XlaOp p,
 
 xla::XlaOp BuildCustomSharding(const xla::XlaOp& input);
 
+xla::XlaOp BuildTpuCustomCall(const xla::XlaOp& x, const xla::XlaOp& y, const std::string& payload);
+
 }  // namespace torch_xla
 
 #endif  // XLA_TORCH_XLA_CSRC_XLA_LOWER_UTIL_H_