Fix special scalar handling for addcdiv and addcmul

Author: ymwangg

test/test_operations_hlo.py

@@ -35,6 +35,21 @@ def test_expand(self):
     hlo_text = torch_xla._XLAC._get_xla_tensors_text([b])
     assert 'aten::expand' in hlo_text
 
+  def test_special_scalars_addcdiv_addcmul(self):
+    a = torch.rand(5, 5).to(xm.xla_device())
+    b = torch.rand(5, 5).to(xm.xla_device())
+    c = torch.rand(5, 5).to(xm.xla_device())
+    for op in [torch.addcdiv, torch.addcmul]:
+      out = op(a, b, c, value=1.0)
+      hlo_text = torch_xla._XLAC._get_xla_tensors_text([out])
+      instructions = hlo_text.split('\n')
+      const_hlo = instructions[1]
+      root_hlo = instructions[5]
+      assert 'prim::Constant()' in const_hlo
+      assert 'xla::device_data()' not in const_hlo
+      assert 'f32' in root_hlo
+      assert 'f64' not in root_hlo
+
 
 if __name__ == '__main__':
   torch.set_default_tensor_type('torch.FloatTensor')

torch_xla/csrc/init_python_bindings.cpp

@@ -741,6 +741,10 @@ absl::flat_hash_map<std::string, absl::variant<int>> ConvertDictToMap(
   return map;
 }
 
+// Override some upstream torch::lazy env vars for better performance.
+// Upstream lazy env vars defined in torch/csrc/lazy/core/config.h.
+void SetDefaultLazyEnvVars() { FLAGS_torch_lazy_handle_special_scalars = true; }
+
 // Maps PT/XLA env vars to upstream torch::lazy env vars.
 // Upstream lazy env vars defined in torch/csrc/lazy/core/config.h.
 void MapXlaEnvVarsToLazy() {
@@ -1469,6 +1473,7 @@ void InitXlaModuleBindings(py::module m) {
         });
 
   m.def("_init_xla_lazy_backend", []() {
+    SetDefaultLazyEnvVars();
     MapXlaEnvVarsToLazy();
     InitXlaBackend();
   });

torch_xla/csrc/xla_lower_util.cpp

@@ -1014,6 +1014,7 @@ xla::XlaOp BuildRoll(xla::XlaOp input, absl::Span<const int64_t> shifts,
 
 xla::XlaOp BuildAddcdiv(xla::XlaOp input, xla::XlaOp t1, xla::XlaOp t2,
                         xla::XlaOp val) {
+  val = MaybeConvertTo(val, XlaHelpers::ShapeOfXlaOp(t1).element_type());
   return XlaHelpers::PromotedAdd(
       input, XlaHelpers::PromotedMul(XlaHelpers::PromotedDiv(t1, t2), val));
 }