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/Patterns/Tensor/Concat.cpp

@@ -2,6 +2,8 @@
 #include "mlir/IR/PatternMatch.h"
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostFoldability.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -18,7 +20,17 @@ struct Concat : public OpConversionPattern<ONNXConcatOp> {
     auto inputs = adaptor.getInputs();
     int64_t axis = adaptor.getAxis();
 
-    rewriter.replaceOp(maxpoolOp, createSpatConcat(rewriter, maxpoolOp.getLoc(), axis, inputs));
+    if (llvm::all_of(inputs, isHostFoldableValue)) {
+      rewriter.replaceOp(maxpoolOp, createSpatConcat(rewriter, maxpoolOp.getLoc(), axis, inputs));
+      return success();
+    }
+
+    auto computeOp = createSpatCompute(
+      rewriter, maxpoolOp.getLoc(), TypeRange {maxpoolOp.getResult().getType()}, {}, inputs, [&](ValueRange args) {
+        spatial::SpatYieldOp::create(
+          rewriter, maxpoolOp.getLoc(), createSpatConcat(rewriter, maxpoolOp.getLoc(), axis, args));
+      });
+    rewriter.replaceOp(maxpoolOp, computeOp.getResults());
 
     return success();
   }

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Gather.cpp

@@ -6,7 +6,7 @@
 #include "llvm/ADT/SmallVector.h"
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Reshape.cpp

@@ -3,7 +3,10 @@
 
 #include "llvm/ADT/SmallVector.h"
 
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostFoldability.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
 using namespace mlir;
@@ -95,18 +98,33 @@ struct Reshape : OpConversionPattern<ONNXReshapeOp> {
       return success();
     }
 
+    auto replaceWithReshape = [&](auto buildReshape) -> LogicalResult {
+      if (isHostFoldableValue(adaptor.getData())) {
+        rewriter.replaceOp(reshapeOp, buildReshape(adaptor.getData()));
+        return success();
+      }
+
+      auto computeOp = createSpatCompute<1>(
+        rewriter, reshapeOp.getLoc(), TypeRange {resultType}, {}, adaptor.getData(), [&](Value data) {
+          Value reshaped = buildReshape(data);
+          spatial::SpatYieldOp::create(rewriter, reshapeOp.getLoc(), reshaped);
+        });
+      rewriter.replaceOp(reshapeOp, computeOp.getResults());
+      return success();
+    };
+
     SmallVector<ReassociationIndices> reassociation;
     if (sourceType.getRank() > resultType.getRank()
-        && inferCollapseReassociation(sourceType.getShape(), resultType.getShape(), reassociation)) {
-      rewriter.replaceOpWithNewOp<tensor::CollapseShapeOp>(reshapeOp, resultType, adaptor.getData(), reassociation);
-      return success();
-    }
+        && inferCollapseReassociation(sourceType.getShape(), resultType.getShape(), reassociation))
+      return replaceWithReshape([&](Value data) {
+        return tensor::CollapseShapeOp::create(rewriter, reshapeOp.getLoc(), resultType, data, reassociation);
+      });
 
     if (sourceType.getRank() < resultType.getRank()
-        && inferExpandReassociation(sourceType.getShape(), resultType.getShape(), reassociation)) {
-      rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(reshapeOp, resultType, adaptor.getData(), reassociation);
-      return success();
-    }
+        && inferExpandReassociation(sourceType.getShape(), resultType.getShape(), reassociation))
+      return replaceWithReshape([&](Value data) {
+        return tensor::ExpandShapeOp::create(rewriter, reshapeOp.getLoc(), resultType, data, reassociation);
+      });
 
     return failure();
   }

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Resize.cpp

@@ -6,7 +6,7 @@
 #include <algorithm>
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Split.cpp

@@ -2,7 +2,9 @@
 #include "mlir/Transforms/DialectConversion.h"
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostFoldability.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
 using namespace mlir;
@@ -47,16 +49,40 @@ struct Split : OpConversionPattern<ONNXSplitOp> {
     outputs.reserve(splitOp.getNumResults());
 
     int64_t offset = 0;
+    SmallVector<RankedTensorType> resultTypes;
+    resultTypes.reserve(splitOp.getNumResults());
+    SmallVector<int64_t> sliceSizes;
+    sliceSizes.reserve(splitOp.getNumResults());
     for (Value result : splitOp.getResults()) {
       auto resultType = dyn_cast<RankedTensorType>(result.getType());
       if (!resultType || !resultType.hasStaticShape())
         return failure();
-      int64_t sliceSize = resultType.getShape()[axis];
-      outputs.push_back(extractSliceAt(adaptor.getInput(), axis, offset, sliceSize, rewriter, splitOp.getLoc()));
-      offset += sliceSize;
+      resultTypes.push_back(resultType);
+      sliceSizes.push_back(resultType.getShape()[axis]);
     }
 
-    rewriter.replaceOp(splitOp, outputs);
+    if (isHostFoldableValue(adaptor.getInput())) {
+      for (int64_t sliceSize : sliceSizes) {
+        outputs.push_back(extractSliceAt(adaptor.getInput(), axis, offset, sliceSize, rewriter, splitOp.getLoc()));
+        offset += sliceSize;
+      }
+      rewriter.replaceOp(splitOp, outputs);
+      return success();
+    }
+
+    auto computeOp = createSpatCompute<1>(
+      rewriter, splitOp.getLoc(), TypeRange(splitOp.getResultTypes()), {}, adaptor.getInput(), [&](Value input) {
+        SmallVector<Value> runtimeOutputs;
+        runtimeOutputs.reserve(resultTypes.size());
+        int64_t runtimeOffset = 0;
+        for (int64_t sliceSize : sliceSizes) {
+          runtimeOutputs.push_back(extractSliceAt(input, axis, runtimeOffset, sliceSize, rewriter, splitOp.getLoc()));
+          runtimeOffset += sliceSize;
+        }
+        spatial::SpatYieldOp::create(rewriter, splitOp.getLoc(), runtimeOutputs);
+      });
+
+    rewriter.replaceOp(splitOp, computeOp.getResults());
     return success();
   }
 };