huge refactor for high RewritePatterns usage and less ad-hoc cpp code

remove Spatial many ops in favor of tensor ops like in pim

src/PIM/Conversion/ONNXToSpatial/Common/Common.hpp

@@ -1,8 +1,7 @@
 #pragma once
 
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Dialect/ONNX/ONNXOps.hpp"
-
 #include "ComputeRegionBuilder.hpp"
 #include "ShapeTilingUtils.hpp"
 #include "WeightMaterialization.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"

src/PIM/Conversion/ONNXToSpatial/Common/ShapeTilingUtils.cpp

@@ -5,6 +5,8 @@
 
 #include "ShapeTilingUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostFoldability.hpp"
 
 using namespace mlir;
 
@@ -30,10 +32,29 @@ SmallVector<Value> sliceTensor(
 
   for (int64_t i = 0; i < numSlices; i++) {
     offsets[axis] = rewriter.getIndexAttr(i * sliceSize);
-    if (i == numSlices - 1 && lastSliceSize != 0)
+    int64_t currentSliceSize = sliceSize;
+    if (i == numSlices - 1 && lastSliceSize != 0) {
+      currentSliceSize = lastSliceSize;
       sizes[axis] = rewriter.getIndexAttr(lastSliceSize);
+    }
 
-    Value slice = tensor::ExtractSliceOp::create(rewriter, loc, tensorToSlice, offsets, sizes, strides);
+    SmallVector<int64_t> sliceShape(shape.begin(), shape.end());
+    sliceShape[axis] = currentSliceSize;
+    auto sliceType =
+      RankedTensorType::get(sliceShape, cast<RankedTensorType>(tensorToSlice.getType()).getElementType());
+
+    Value slice;
+    if (isHostFoldableValue(tensorToSlice)) {
+      slice = tensor::ExtractSliceOp::create(rewriter, loc, tensorToSlice, offsets, sizes, strides);
+    }
+    else {
+      auto sliceCompute =
+        createSpatCompute<1>(rewriter, loc, TypeRange {sliceType}, {}, ValueRange {tensorToSlice}, [&](Value input) {
+          Value computedSlice = tensor::ExtractSliceOp::create(rewriter, loc, input, offsets, sizes, strides);
+          spatial::SpatYieldOp::create(rewriter, loc, computedSlice);
+        });
+      slice = sliceCompute.getResult(0);
+    }
     slices.push_back(slice);
   }
 

src/PIM/Conversion/ONNXToSpatial/Common/ShapeTilingUtils.hpp

@@ -5,15 +5,15 @@
 #include "mlir/IR/Value.h"
 #include "mlir/Transforms/DialectConversion.h"
 
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+
 #include <cassert>
 #include <cstddef>
 #include <type_traits>
 #include <utility>
 
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/SmallVector.h"
-
 namespace onnx_mlir {
 
 template <class ShapedType>
@@ -105,7 +105,8 @@ inline auto getTensorShape(mlir::Value tensor) {
 inline bool haveSameStaticShape(mlir::Value lhs, mlir::Value rhs) {
   auto lhsType = mlir::dyn_cast<mlir::RankedTensorType>(lhs.getType());
   auto rhsType = mlir::dyn_cast<mlir::RankedTensorType>(rhs.getType());
-  return lhsType && rhsType && lhsType.hasStaticShape() && rhsType.hasStaticShape() && lhsType.getShape() == rhsType.getShape();
+  return lhsType && rhsType && lhsType.hasStaticShape() && rhsType.hasStaticShape()
+      && lhsType.getShape() == rhsType.getShape();
 }
 
 /// Slices a statically shaped tensor along one axis into contiguous pieces of

src/PIM/Conversion/ONNXToSpatial/Common/WeightMaterialization.cpp

@@ -5,12 +5,12 @@
 #include "mlir/IR/Value.h"
 #include "mlir/Support/LogicalResult.h"
 
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/ADT/STLExtras.h"
 
-#include "WeightMaterialization.hpp"
 #include "ShapeTilingUtils.hpp"
+#include "WeightMaterialization.hpp"
 #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -28,7 +28,7 @@ bool isWeightLikeComputeOperand(Value value) {
   while (auto* definingOp = value.getDefiningOp()) {
     if (!visited.insert(definingOp).second)
       return false;
-    if (hasWeightAlways(definingOp))
+    if (isa<arith::ConstantOp, ONNXConstantOp>(definingOp) || hasWeightAlways(definingOp))
       return true;
 
     if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(definingOp)) {