Fix reflected elementwise binary ops (#68)

Author: tkarna

examples/shallow_water.py

@@ -233,7 +233,7 @@ def rhs(u, v, e):
         # potential vorticity
         dudy[:, 1:-1] = (u[:, 1:] - u[:, :-1]) / dy
         dvdx[1:-1, :] = (v[1:, :] - v[:-1, :]) / dx
-        q[:, :] = (dvdx - dudy + coriolis) / H_at_f
+        q[:, :] = (coriolis - dudy + dvdx) / H_at_f
 
         # Advection of potential vorticity, Arakawa and Hsu (1990)
         # Define alpha, beta, gamma, delta for each cell in T points

src/EWBinOp.cpp

@@ -84,6 +84,26 @@ static ::imex::ptensor::EWBinOpId ddpt2mlir(const EWBinOpId bop) {
   }
 }
 
+bool is_reflected_op(EWBinOpId op) {
+  switch (op) {
+  case __RADD__:
+  case __RAND__:
+  case __RFLOORDIV__:
+  case __RLSHIFT__:
+  case __RMOD__:
+  case __RMUL__:
+  case __ROR__:
+  case __RPOW__:
+  case __RRSHIFT__:
+  case __RSUB__:
+  case __RTRUEDIV__:
+  case __RXOR__:
+    return true;
+  default:
+    return false;
+  }
+}
+
 struct DeferredEWBinOp : public Deferred {
   id_type _a;
   id_type _b;
@@ -107,8 +127,16 @@ struct DeferredEWBinOp : public Deferred {
         ::imex::ptensor::toMLIR(builder, DDPT::jit::getPTDType(_dtype));
     auto outTyp = aTyp.cloneWith(shape(), outElemType);
 
+    ::mlir::Value one, two;
+    if (is_reflected_op(_op)) {
+      one = bv;
+      two = av;
+    } else {
+      one = av;
+      two = bv;
+    }
     auto bop = builder.create<::imex::ptensor::EWBinOp>(
-        loc, outTyp, builder.getI32IntegerAttr(ddpt2mlir(_op)), av, bv);
+        loc, outTyp, builder.getI32IntegerAttr(ddpt2mlir(_op)), one, two);
 
     dm.addVal(this->guid(), bop,
               [this](uint64_t rank, void *l_allocated, void *l_aligned,

test/test_ewb.py

@@ -182,3 +182,17 @@ def test_add_typecast(self):
             assert c.dtype == ctype
             c2 = dt.to_numpy(c)
             assert numpy.allclose(c2, [1, 2, 3, 4, 5, 6, 7, 8])
+
+    def test_reflected_sub(self):
+        a = dt.full((4,), 4, dtype=dt.float64)
+        b = 10.0
+        c = b - a
+        c2 = dt.to_numpy(c)
+        assert numpy.allclose(c2, [6, 6, 6, 6])
+
+    def test_reflected_div(self):
+        a = dt.full((4,), 2, dtype=dt.float64)
+        b = 10.0
+        c = b / a
+        c2 = dt.to_numpy(c)
+        assert numpy.allclose(c2, [5, 5, 5, 5])