Use promoteAdd/Div/Mul

Author: JackCaoG

torch_xla/csrc/ops/ops_lower_fn.cpp

@@ -7,6 +7,7 @@
 #include "torch_xla/csrc/matrix.h"
 #include "torch_xla/csrc/pooling.h"
 #include "torch_xla/csrc/reduction.h"
+#include "torch_xla/csrc/xla_lower_util.h"
 
 namespace torch_xla {
 
@@ -63,7 +64,7 @@ torch_xla::XlaOpVector Addcdiv::Lower(LoweringContext* loctx) const {
   xla::XlaOp xla_t1 = loctx->GetOutputOp(operand(1));
   xla::XlaOp xla_t2 = loctx->GetOutputOp(operand(2));
   xla::XlaOp xla_val = loctx->GetOutputOp(operand(3));
-  return ReturnOp(xla_input + (xla_t1 / xla_t2) * xla_val, loctx);
+  return ReturnOp(BuildAddcdiv(xla_input, xla_t1, xla_t2, xla_val), loctx);
 }
 
 torch_xla::XlaOpVector Addcmul::Lower(LoweringContext* loctx) const {
@@ -75,7 +76,7 @@ torch_xla::XlaOpVector Addcmul::Lower(LoweringContext* loctx) const {
   xla::XlaOp xla_t1 = loctx->GetOutputOp(operand(1));
   xla::XlaOp xla_t2 = loctx->GetOutputOp(operand(2));
   xla::XlaOp xla_val = loctx->GetOutputOp(operand(3));
-  return ReturnOp(xla_input + (xla_t1 * xla_t2) * xla_val, loctx);
+  return ReturnOp(BuildAddcmul(xla_input, xla_t1, xla_t2, xla_val), loctx);
 }
 
 torch_xla::XlaOpVector Asin::Lower(LoweringContext* loctx) const {

torch_xla/csrc/ops/ops_xla_shape_fn.cpp

@@ -7,6 +7,7 @@
 #include "torch_xla/csrc/helpers.h"
 #include "torch_xla/csrc/pooling.h"
 #include "torch_xla/csrc/reduction.h"
+#include "torch_xla/csrc/xla_lower_util.h"
 
 namespace torch_xla {
 namespace {
@@ -85,7 +86,7 @@ xla::Shape AddcdivOutputShape(const torch::lazy::Value& input,
                               const torch::lazy::Value& t2,
                               const torch::lazy::Value& value) {
   auto shape_fn = [](absl::Span<const xla::XlaOp> operands) -> xla::XlaOp {
-    return operands[0] + (operands[1] / operands[2]) * operands[3];
+    return BuildAddcdiv(operands[0], operands[1], operands[2], operands[3]);
   };
   return InferOutputShape({GetXlaShape(input), GetXlaShape(t1), GetXlaShape(t2),
                            GetXlaShape(value)},
@@ -97,8 +98,9 @@ xla::Shape AddcmulOutputShape(const torch::lazy::Value& input,
                               const torch::lazy::Value& t2,
                               const torch::lazy::Value& value) {
   auto shape_fn = [](absl::Span<const xla::XlaOp> operands) -> xla::XlaOp {
-    return operands[0] + (operands[1] * operands[2]) * operands[3];
+    return BuildAddcmul(operands[0], operands[1], operands[2], operands[3]);
   };
+
   return InferOutputShape({GetXlaShape(input), GetXlaShape(t1), GetXlaShape(t2),
                            GetXlaShape(value)},
                           shape_fn);

torch_xla/csrc/xla_lower_util.cpp

@@ -1018,4 +1018,16 @@ xla::XlaOp BuildRoll(xla::XlaOp input, absl::Span<const int64_t> shifts,
   return need_flatten ? xla::Reshape(input, input_shape.dimensions()) : input;
 }
 
+xla::XlaOp BuildAddcdiv(xla::XlaOp input, xla::XlaOp t1, xla::XlaOp t2,
+                        xla::XlaOp val) {
+  return XlaHelpers::PromotedAdd(
+      input, XlaHelpers::PromotedMul(XlaHelpers::PromotedDiv(t1, t2), val));
+}
+
+xla::XlaOp BuildAddcmul(xla::XlaOp input, xla::XlaOp t1, xla::XlaOp t2,
+                        xla::XlaOp val) {
+  return XlaHelpers::PromotedAdd(
+      input, XlaHelpers::PromotedMul(XlaHelpers::PromotedMul(t1, t2), val));
+}
+
 }  // namespace torch_xla

torch_xla/csrc/xla_lower_util.h

@@ -121,4 +121,10 @@ xla::XlaOp BuildXLogY(xla::XlaOp input, xla::XlaOp other);
 xla::XlaOp BuildRoll(xla::XlaOp input, absl::Span<const int64_t> shifts,
                      absl::Span<const int64_t> dims);
 
+xla::XlaOp BuildAddcdiv(xla::XlaOp input, xla::XlaOp t1, xla::XlaOp t2,
+                        xla::XlaOp val);
+
+xla::XlaOp BuildAddcmul(xla::XlaOp input, xla::XlaOp t1, xla::XlaOp t2,
+                        xla::XlaOp val);
+
 }  // namespace torch_xla