Reshape Reloaded (#2)

- resurrecting always-copying reshape
- adjusting to fixes in imex (dtype-independent wait, async)
- enabling tests using reshape
- temporarily disabling async execution of reshape until fixed in IMEX

Author: fschlimb

.github/workflows/ci.yml

@@ -1,4 +1,4 @@
-name: GitHub Actions Demo
+name: Sharpy CI
 run-name: ${{ github.actor }} CI for sharpy
 on:
   push:
@@ -41,7 +41,8 @@ jobs:
           rm -rf $CONDA_ROOT
           cd $GITHUB_WORKSPACE/..
           rm -f Miniconda3-*.sh
-          CPKG=Miniconda3-latest-Linux-x86_64.sh
+          # CPKG=Miniconda3-latest-Linux-x86_64.sh
+          CPKG=Miniconda3-py311_24.3.0-0-Linux-x86_64.sh
           wget -q https://repo.anaconda.com/miniconda/$CPKG
           bash $CPKG -u -b -f -p $CONDA_ROOT
           export PATH=$CONDA_ROOT/condabin:$CONDA_ROOT/bin:${PATH}

conda-recipe/build.sh

@@ -53,13 +53,12 @@ if [ ! -d "${INSTALLED_DIR}/imex/lib" ]; then
     rm -rf ${INSTALLED_DIR}/imex
     IMEX_SHA=$(cat imex_version.txt)
     if [ ! -d "mlir-extensions" ]; then
-        git clone --recurse-submodules --branch main --single-branch https://github.com/intel/mlir-extensions
+        git clone --recurse-submodules https://github.com/intel/mlir-extensions
     fi
     pushd mlir-extensions
     git reset --hard HEAD
     git fetch --prune
     git checkout $IMEX_SHA
-    git apply ${RECIPE_DIR}/imex_*.patch
     LLVM_SHA=$(cat build_tools/llvm_version.txt)
     # if [ ! -d "llvm-project" ]; then ln -s ~/github/llvm-project .; fi
     if [ ! -d "llvm-project" ]; then

conda-recipe/imex_findsycl.patch

@@ -197,8 +197,8 @@ def main():
     # * Analyze and output results.
     # ******************************************************************************
 
-    B = np.spmd.gather(B)
-    norm = np.linalg.norm(np.reshape(B, n * n), ord=1)
+    B = np.spmd.gather(np.reshape(B, (n * n,)))
+    norm = np.linalg.norm(B, ord=1)
     active_points = (n - 2 * r) ** 2
     norm /= active_points
 

examples/stencil-2d.py

@@ -1 +1 @@
-617949ac6105f28faeab4fa3018142195d1125c0
+a6109b1005932d8b4c1d2e8ab0ec4abe7411762a

imex_version.txt

@@ -43,7 +43,7 @@ def build_cmake(self, ext):
         os.chdir(str(build_temp))
         self.spawn(["cmake", str(cwd)] + cmake_args)
         if not self.dry_run:
-            self.spawn(["cmake", "--build", "."] + build_args)
+            self.spawn(["cmake", "--build", ".", "-j5"] + build_args)
         # Troubleshooting: if fail on line above then delete all possible
         # temporary CMake files including "CMakeCache.txt" in top level dir.
         os.chdir(str(cwd))

setup.py

@@ -107,6 +107,14 @@ def _validate_device(device):
             f"{func} = lambda start, end, step, endpoint, dtype=float64, device='', team=1: ndarray(_csp.Creator.linspace(start, end, step, endpoint, dtype, _validate_device(device), team))"
         )
 
+
+for func in api.api_categories["ManipOp"]:
+    FUNC = func.upper()
+    if func == "reshape":
+        exec(
+            f"{func} = lambda this, shape, cp=None: ndarray(_csp.ManipOp.reshape(this._t, shape, cp))"
+        )
+
 for func in api.api_categories["ReduceOp"]:
     FUNC = func.upper()
     exec(

sharpy/__init__.py

@@ -175,7 +175,7 @@
             "concat",  # (arrays, /, *, axis=0)
             "expand_dims",  # (x, /, *, axis)
             "flip",  # (x, /, *, axis=None)
-            "reshape",  # (x, /, shape)
+            "reshape",  # (x, /, shape, *, copy: bool | None = None)
             "roll",  # (x, /, shift, *, axis=None)
             "squeeze",  # (x, /, axis)
             "stack",  # (arrays, /, *, axis=0)

sharpy/array_api.py

@@ -12,6 +12,7 @@
 #include "sharpy/jit/mlir.hpp"
 
 #include <imex/Dialect/Dist/IR/DistOps.h>
+#include <imex/Dialect/Dist/Utils/Utils.h>
 #include <imex/Dialect/NDArray/IR/NDArrayOps.h>
 #include <mlir/IR/Builders.h>
 
@@ -41,7 +42,7 @@ struct DeferredReshape : public Deferred {
             : ::imex::getIntAttr(builder, COPY_ALWAYS ? true : false, 1);
 
     auto aTyp = av.getType().cast<::imex::ndarray::NDArrayType>();
-    auto outTyp = aTyp.cloneWith(shape(), aTyp.getElementType());
+    auto outTyp = imex::dist::cloneWithShape(aTyp, shape());
 
     auto op =
         builder.create<::imex::ndarray::ReshapeOp>(loc, outTyp, av, shp, copyA);