Elementwise binop type promotion (#67)

Author: tkarna

examples/shallow_water.py

@@ -75,11 +75,10 @@ def info(s):
     coriolis = 10.0
 
     # domain extent
-    # NOTE need to be floats
-    xmin = -1.0
-    xmax = 1.0
-    ymin = -1.0
-    ymax = 1.0
+    xmin = -1
+    xmax = 1
+    ymin = -1
+    ymax = 1
     lx = xmax - xmin
     ly = ymax - ymin
 

examples/wave_equation.py

@@ -75,11 +75,10 @@ def info(s):
     g = 9.81
 
     # domain extent
-    # NOTE need to be floats
-    xmin = -1.0
-    xmax = 1.0
-    ymin = -1.0
-    ymax = 1.0
+    xmin = -1
+    xmax = 1
+    ymin = -1
+    ymax = 1
     lx = xmax - xmin
     ly = ymax - ymin
 

imex_version.txt

@@ -1 +1 @@
-dd921a5893d2956ddee3d7fecb84612edf15fbbe
+3067b8fb5a239b46a4a563c200c7fb8abb066643

src/EWBinOp.cpp

@@ -92,18 +92,20 @@ struct DeferredEWBinOp : public Deferred {
   DeferredEWBinOp() = default;
   DeferredEWBinOp(EWBinOpId op, const tensor_i::future_type &a,
                   const tensor_i::future_type &b)
-      : Deferred(a.dtype(), broadcast(a.shape(), b.shape()), a.device(),
-                 a.team()),
+      : Deferred(promoted_dtype(a.dtype(), b.dtype()),
+                 broadcast(a.shape(), b.shape()), a.device(), a.team()),
         _a(a.guid()), _b(b.guid()), _op(op) {}
 
   bool generate_mlir(::mlir::OpBuilder &builder, const ::mlir::Location &loc,
                      jit::DepManager &dm) override {
-    // FIXME the type of the result is based on a only
     auto av = dm.getDependent(builder, _a);
     auto bv = dm.getDependent(builder, _b);
 
     auto aTyp = av.getType().cast<::imex::ptensor::PTensorType>();
-    auto outTyp = aTyp.cloneWith(shape(), aTyp.getElementType());
+    auto bTyp = bv.getType().cast<::imex::ptensor::PTensorType>();
+    auto outElemType =
+        ::imex::ptensor::toMLIR(builder, DDPT::jit::getPTDType(_dtype));
+    auto outTyp = aTyp.cloneWith(shape(), outElemType);
 
     auto bop = builder.create<::imex::ptensor::EWBinOp>(
         loc, outTyp, builder.getI32IntegerAttr(ddpt2mlir(_op)), av, bv);

src/include/ddptensor/CppTypes.hpp

@@ -156,6 +156,113 @@ static size_t sizeof_dtype(const DTypeId dt) {
   };
 };
 
+static bool is_float(DTypeId t) {
+  switch (t) {
+  case DDPT::DTypeId::FLOAT64:
+  case DDPT::DTypeId::FLOAT32:
+    return true;
+  default:
+    return false;
+  }
+}
+
+static bool is_int(DTypeId t) {
+  switch (t) {
+  case DDPT::DTypeId::INT64:
+  case DDPT::DTypeId::INT32:
+  case DDPT::DTypeId::INT16:
+  case DDPT::DTypeId::INT8:
+    return true;
+  default:
+    return false;
+  }
+}
+
+static bool is_uint(DTypeId t) {
+  switch (t) {
+  case DDPT::DTypeId::UINT64:
+  case DDPT::DTypeId::UINT32:
+  case DDPT::DTypeId::UINT16:
+  case DDPT::DTypeId::UINT8:
+  case DDPT::DTypeId::BOOL:
+    return true;
+  default:
+    return false;
+  }
+}
+
+static size_t dtype_bitwidth(DTypeId t) {
+  switch (t) {
+  case DDPT::DTypeId::FLOAT64:
+  case DDPT::DTypeId::INT64:
+  case DDPT::DTypeId::UINT64:
+    return 64;
+  case DDPT::DTypeId::FLOAT32:
+  case DDPT::DTypeId::INT32:
+  case DDPT::DTypeId::UINT32:
+    return 32;
+  case DDPT::DTypeId::INT16:
+  case DDPT::DTypeId::UINT16:
+    return 16;
+  case DDPT::DTypeId::INT8:
+  case DDPT::DTypeId::UINT8:
+    return 8;
+  case DDPT::DTypeId::BOOL:
+    return 1;
+  default:
+    assert(!"Unknown DTypeId");
+  }
+}
+
+static DTypeId get_float_dtype(size_t bitwidth) {
+  switch (bitwidth) {
+  case 64:
+    return DDPT::DTypeId::FLOAT64;
+  case 32:
+    return DDPT::DTypeId::FLOAT32;
+  default:
+    assert(!"Unknown bitwidth");
+  }
+}
+
+static DTypeId get_int_dtype(size_t bitwidth) {
+  switch (bitwidth) {
+  case 64:
+    return DDPT::DTypeId::INT64;
+  case 32:
+    return DDPT::DTypeId::INT32;
+  case 16:
+    return DDPT::DTypeId::INT16;
+  case 8:
+    return DDPT::DTypeId::INT8;
+  default:
+    assert(!"Unknown bitwidth");
+  }
+}
+
+static DTypeId promoted_dtype(DTypeId a, DTypeId b) {
+  if ((is_float(a) && is_float(b)) || (is_int(a) && is_int(b)) ||
+      (is_uint(a) && is_uint(b))) {
+    return dtype_bitwidth(a) > dtype_bitwidth(b) ? a : b;
+  }
+  if (is_float(a) || is_float(b)) {
+    return get_float_dtype(std::max(dtype_bitwidth(a), dtype_bitwidth(b)));
+  }
+  // mixed signed/unsigned int case
+  size_t si_width, ui_width, max_width = 64;
+  if (is_uint(a)) {
+    ui_width = dtype_bitwidth(a);
+    si_width = dtype_bitwidth(b);
+  } else {
+    ui_width = dtype_bitwidth(b);
+    si_width = dtype_bitwidth(a);
+  }
+  if (ui_width < si_width) {
+    return get_int_dtype(si_width);
+  }
+  return get_int_dtype(std::min(2 * ui_width, max_width));
+}
+
 using RedOpType = ReduceOpId;
 
 inline RedOpType red_op(const char *op) {

src/jit/mlir.cpp

@@ -135,7 +135,7 @@ ::mlir::SmallVector<::mlir::Attribute> mkEnvs(::mlir::Builder &builder,
   }
   if (!device.empty()) {
     envs.emplace_back(
-        ::imex::ptensor::GPUEnvAttr::get(builder.getStringAttr(device)));
+        ::imex::region::GPUEnvAttr::get(builder.getStringAttr(device)));
   }
   return envs;
 }

test/test_ewb.py

@@ -1,4 +1,5 @@
 import ddptensor as dt
+import numpy
 from utils import runAndCompare, mpi_dtypes, on_gpu
 import pytest
 
@@ -162,3 +163,22 @@ def test_bitwise_and(self):
             r1 = dt.sum(c, [0, 1])
             v = 6 * 6 * 2
             assert float(r1) == v
+
+    def test_add_typecast(self):
+        type_list = [
+            (dt.int8, dt.int64, dt.int64),
+            (dt.int64, dt.int8, dt.int64),
+            (dt.int8, dt.float64, dt.float64),
+            (dt.float64, dt.int8, dt.float64),
+            (dt.float32, dt.float64, dt.float64),
+            (dt.float32, dt.int64, dt.float64),
+            (dt.int32, dt.uint32, dt.int64),
+            (dt.int32, dt.uint64, dt.int64),
+        ]
+        for atype, btype, ctype in type_list:
+            a = dt.arange(0, 8, 1, dtype=atype)
+            b = dt.ones((8,), dtype=btype)
+            c = a + b
+            assert c.dtype == ctype
+            c2 = dt.to_numpy(c)
+            assert numpy.allclose(c2, [1, 2, 3, 4, 5, 6, 7, 8])