adding examples (#22)

* adding examples and a new test
* delete -> free; new -> aligned_alloc
* updating imex

---------

Co-authored-by: Tuomas Karna <tuomas.karna@intel.com>

Author: fschlimb

.github/workflows/ci.yml

@@ -122,18 +122,28 @@ jobs:
           . $GITHUB_WORKSPACE/third_party/install/miniconda/etc/profile.d/conda.sh
           . $GITHUB_WORKSPACE/third_party/install/miniconda/bin/activate ddpt
           MLIRROOT=$GITHUB_WORKSPACE/third_party/install/llvm-mlir IMEXROOT=$GITHUB_WORKSPACE/third_party/install/imex python setup.py install
+          DDPT_ROOT=`pip show -f ddptensor | grep Location | awk '{print $2}'`
+          echo "DDPT_IDTR_SO=${DDPT_ROOT}/ddptensor/libidtr.so" >> $GITHUB_ENV
+          echo "DDPT_CRUNNER_SO=${GITHUB_WORKSPACE}/third_party/install/llvm-mlir/lib/libmlir_c_runner_utils.so" >> $GITHUB_ENV
       - name: Test DDPT (pytest)
         run: |
           . $GITHUB_WORKSPACE/third_party/install/miniconda/etc/profile.d/conda.sh
           . $GITHUB_WORKSPACE/third_party/install/miniconda/bin/activate ddpt
-          DDPT_ROOT=`pip show -f ddptensor | grep Location | awk '{print $2}'`
-          export DDPT_IDTR_SO=${DDPT_ROOT}/ddptensor/libidtr.so
-          export DDPT_CRUNNER_SO="$GITHUB_WORKSPACE"/third_party/install/llvm-mlir/lib/libmlir_c_runner_utils.so
-          pytest test
-          DDPT_FORCE_DIST=1 pytest test
+          cd test
+          pytest .
+          DDPT_FORCE_DIST=1 pytest .
+          cd -
+      - name: Run examples
+        run: |
+          . $GITHUB_WORKSPACE/third_party/install/miniconda/etc/profile.d/conda.sh
+          . $GITHUB_WORKSPACE/third_party/install/miniconda/bin/activate ddpt
+          cd examples
+          python -u ./stencil-2d.py 5 1024 star 4
+          DDPT_FORCE_DIST=1 python -u ./stencil-2d.py 5 1024 star 4
+          python -u ./wave_equation.py -ct
+          DDPT_FORCE_DIST=1 python -u ./wave_equation.py -ct
+          cd -
       - name: Cleanup
         run: |
           pip list
           pip uninstall -y ddptensor
-      - run: |
-          echo "This job's status is ${{ job.status }}."

imex_version.txt

@@ -1 +1 @@
-399d083586cac8b2cb71f7a8da514ac2d816aac4
+865cea0c1af5fe7ac1458907637c28f69afd7146

setup.py

@@ -33,7 +33,7 @@ def build_cmake(self, ext):
             "-DCMAKE_BUILD_TYPE=" + config,
             "-DCMAKE_VERBOSE_MAKEFILE=ON",
             "-G=Ninja",
-            "-DLLVM_ENABLE_LLD=ON"
+            "-DLLVM_ENABLE_LLD=ON",
         ]
 
         # example of build args
@@ -56,7 +56,7 @@ def build_cmake(self, ext):
     name="ddptensor",
     version="0.1",
     description="Distributed Tensor and more",
-    packages=["ddptensor"],  # , "ddptensor.numpy", "ddptensor.torch"],
+    packages=["ddptensor", "ddptensor.numpy"],  # "ddptensor.torch"],
     ext_modules=[CMakeExtension("ddptensor/_ddptensor")],
     cmdclass=dict(
         # Enable the CMakeExtension entries defined above

src/DDPTensorImpl.cpp

@@ -58,10 +58,10 @@ DDPTensorImpl::DDPTensorImpl(const int64_t *shape, uint64_t N, rank_type owner)
 
 DDPTensorImpl::~DDPTensorImpl() {
   if (!_base) {
-    delete[] reinterpret_cast<char *>(_allocated);
+    free(_allocated); // delete[] reinterpret_cast<char *>(_allocated);
   }
-  delete[] _gs_allocated;
-  delete[] _lo_allocated;
+  free(_gs_allocated);
+  free(_lo_allocated);
   delete[] _sizes;
   delete[] _strides;
 }
@@ -72,7 +72,8 @@ DDPTensorImpl::ptr_type DDPTensorImpl::clone(bool copy) {
   auto sz = size();
   auto esz = sizeof_dtype(dtype());
   auto bsz = sz * esz;
-  auto allocated = new (std::align_val_t(esz)) char[bsz];
+  auto allocated =
+      aligned_alloc(esz, bsz); // new (std::align_val_t(esz)) char[bsz];
   auto aligned = allocated;
   if (copy)
     memcpy(aligned, _aligned, bsz);
@@ -82,7 +83,8 @@ DDPTensorImpl::ptr_type DDPTensorImpl::clone(bool copy) {
   // auto gs_aligned = gs_allocated;
   auto gs_allocated = _gs_allocated;
   auto gs_aligned = _gs_aligned;
-  auto lo_allocated = new uint64_t[nd];
+  auto lo_allocated =
+      (uint64_t *)aligned_alloc(sizeof(uint64_t), nd * sizeof(uint64_t));
   auto lo_aligned = lo_allocated;
   memcpy(lo_aligned, _lo_aligned, nd * sizeof(*lo_aligned));
 
@@ -94,7 +96,9 @@ DDPTensorImpl::ptr_type DDPTensorImpl::clone(bool copy) {
 
 void DDPTensorImpl::alloc(bool all) {
   auto esz = sizeof_dtype(_dtype);
-  _allocated = new (std::align_val_t(esz)) char[esz * local_size()];
+  _allocated =
+      aligned_alloc(esz, esz * local_size()); // new (std::align_val_t(esz))
+                                              // char[esz * local_size()];
   _aligned = _allocated;
   _offset = 0;
   if (all) {

test/test_temparray.py

@@ -0,0 +1,46 @@
+import ddptensor as dt
+import pytest
+
+
+class TestTemporaries:
+    def test_temp_inline(self):
+        dtyp = dt.float64
+        a = dt.ones((6, 6), dtype=dtyp)
+        b = dt.ones((6, 6), dtype=dtyp)
+
+        dt.sync()
+        c = a[0:5, 0:5]
+        b[0:5, 0:5] = b[0:5, 0:5] + c
+        del c
+        dt.sync()
+
+        c = a[0:5, 0:5]
+        b[0:5, 0:5] = b[0:5, 0:5] + c
+        del c
+        dt.sync()
+
+        r1 = dt.sum(b, [0, 1])
+        v = 5 * 5 * 3 + 6 + 5
+        assert float(r1) == v
+
+    def test_temp_func(self):
+        def func(a):
+            return a[0:5, 0:5]
+
+        def update(a, b):
+            c = func(a)
+            b[0:5, 0:5] = b[0:5, 0:5] + c
+
+        dtyp = dt.float64
+        a = dt.ones((6, 6), dtype=dtyp)
+        b = dt.ones((6, 6), dtype=dtyp)
+
+        dt.sync()
+        update(a, b)
+        dt.sync()
+        update(a, b)
+        dt.sync()
+
+        r1 = dt.sum(b, [0, 1])
+        v = 5 * 5 * 3 + 6 + 5
+        assert float(r1) == v

test/utils.py

@@ -7,9 +7,6 @@ def runAndCompare(func, do_gather=True):
     a = ddptensor.spmd.gather(aa) if do_gather else aa
     b = func(numpy)
     if isinstance(b, numpy.ndarray):
-        print(aa)
-        print(a)
-        print(b)
         return a.shape == b.shape and numpy.allclose(a, b, rtol=1e-8, atol=1e-8)
     return float(a) == float(b)