Full codegen addcmul

pytorch · May 13, 2022 · 66e0da5 · 66e0da5
1 parent 56a52d5
commit 66e0da5
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 2 deletions.
diff --git a/torch_xla/csrc/ops/ops_lower_fn.cpp b/torch_xla/csrc/ops/ops_lower_fn.cpp
@@ -21,6 +21,15 @@ torch_xla::XlaOpVector Acosh::Lower(LoweringContext* loctx) const {
   return ReturnOp(xla::Acosh(xla_input), loctx);
 }
 
+torch_xla::XlaOpVector Addcmul::Lower(LoweringContext* loctx) const {
+  xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
+  xla::XlaOp xla_tensor_1 = loctx->GetOutputOp(operand(1));
+  xla::XlaOp xla_tensor_2 = loctx->GetOutputOp(operand(2));
+  xla::XlaOp xla_constant = loctx->GetOutputOp(operand(3));
+  xla::XlaOp mul = xla_tensor_1 * xla_tensor_2;
+  return ReturnOp(xla_input + mul * xla_constant, loctx);
+}
+
 torch_xla::XlaOpVector Asin::Lower(LoweringContext* loctx) const {
   xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
   return ReturnOp(xla::Asin(xla_input), loctx);

diff --git a/torch_xla/csrc/ops/ops_xla_shape_fn.cpp b/torch_xla/csrc/ops/ops_xla_shape_fn.cpp
@@ -10,6 +10,8 @@ xla::Shape AcosOutputShape(const XlaValue& input) { return input.xla_shape(); }
 
 xla::Shape AcoshOutputShape(const XlaValue& input) { return input.xla_shape(); }
 
+xla::Shape AddcmulOutputShape(const XlaValue& input, const XlaValue& tensor1, const XlaValue& tensor2, const XlaValue& value) { return input.xla_shape(); }
+
 xla::Shape AsinOutputShape(const XlaValue& input) { return input.xla_shape(); }
 
 xla::Shape AsinhOutputShape(const XlaValue& input) { return input.xla_shape(); }

diff --git a/torch_xla/csrc/ops/ops_xla_shape_fn.h b/torch_xla/csrc/ops/ops_xla_shape_fn.h
@@ -9,6 +9,8 @@ xla::Shape AcosOutputShape(const XlaValue& input);
 
 xla::Shape AcoshOutputShape(const XlaValue& input);
 
+xla::Shape AddcmulOutputShape(const XlaValue& input, const XlaValue& tensor1, const XlaValue& tensor2, const XlaValue& value);
+
 xla::Shape AsinOutputShape(const XlaValue& input);
 
 xla::Shape AsinhOutputShape(const XlaValue& input);

diff --git a/xla_native_functions.yaml b/xla_native_functions.yaml
@@ -1,9 +1,10 @@
 backend: XLA
 cpp_namespace: torch_xla
 full_codegen:
+  - abs
   - acos
   - acosh
-  - abs
+  - addcmul
   - asin
   - asinh
   - atan
@@ -42,7 +43,6 @@ supported:
   - add.Tensor
   - addcdiv
   - addcdiv_
-  - addcmul
   - addmm
   - alias
   - all