allow ':' slice (#24)

* adding support for slice ':'
* adding tests
* updating imex sha

Author: fschlimb

imex_version.txt

@@ -1 +1 @@
-e9b8ef1bb2126a4aa544ae1b25eae33f7f9a2d80
+d944080e43cc67798b92f1ca5d1ea8e50df69eb5

src/SetGetItem.cpp

@@ -19,6 +19,7 @@
 #include <imex/Dialect/Dist/IR/DistOps.h>
 #include <imex/Dialect/PTensor/IR/PTensorOps.h>
 #include <imex/Utils/PassUtils.h>
+#include <mlir/Dialect/Linalg/Utils/Utils.h>
 #include <mlir/IR/Builders.h>
 
 #include <pybind11/numpy.h>
@@ -168,7 +169,12 @@ struct DeferredSetItem : public Deferred {
     std::vector<::mlir::Value> stridesV(nd);
     for (auto i = 0; i < nd; ++i) {
       offsV[i] = ::imex::createIndex(loc, builder, offs[i]);
-      sizesV[i] = ::imex::createIndex(loc, builder, sizes[i]);
+      if (sizes[i] == ALL_SIZE) {
+        sizesV[i] =
+            builder.create<::imex::ptensor::DimOp>(loc, av, i).getResult();
+      } else {
+        sizesV[i] = ::imex::createIndex(loc, builder, sizes[i]);
+      }
       stridesV[i] = ::imex::createIndex(loc, builder, strides[i]);
     }
     // insertsliceop has no return value, so we just create the op...
@@ -280,7 +286,12 @@ struct DeferredGetItem : public Deferred {
         sizesV[i] = builder.getIndexAttr(sizes[i]);
         shape[i] = sizes[i];
       } else {
-        sizesV[i] = ::imex::createIndex(loc, builder, sizes[i]);
+        if (sizes[i] == ALL_SIZE) {
+          sizesV[i] =
+              builder.create<::imex::ptensor::DimOp>(loc, av, i).getResult();
+        } else {
+          sizesV[i] = ::imex::createIndex(loc, builder, sizes[i]);
+        }
       }
     }
 

test/test_ewb.py

@@ -13,6 +13,14 @@
 
 
 class TestEWB:
+    def test_add0(self):
+        for dtyp in mpi_dtypes:
+            a = dt.ones((6, 6), dtype=dtyp)
+            c = a + a
+            r1 = dt.sum(c)
+            v = 6 * 6 * 2
+            assert float(r1) == v
+
     def test_add1(self):
         for dtyp in mpi_dtypes:
             a = dt.ones((6, 6), dtype=dtyp)

test/test_setget.py

@@ -1,3 +1,4 @@
+import ddptensor as dt
 from utils import runAndCompare
 import pytest
 
@@ -39,3 +40,10 @@ def doit(aapi):
             return a
 
         assert runAndCompare(doit)
+
+    def test_colon(self):
+        a = dt.ones((16, 16), dt.float64)
+        b = dt.zeros((16, 16), dt.float64)
+        a[:, :] = b[:, :]
+        r1 = dt.sum(a)
+        assert float(r1) == 0