use affine dialect to express simple constant progressions

run dce at the end of MaterializeMergeSchedule to get rid of unused constants

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -1,3 +1,5 @@
+#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
@@ -143,7 +145,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   OperationFolder constantFolder(&getContext());
 
   ConversionTarget target(*ctx);
-  target.addLegalDialect<PimDialect,
+  target.addLegalDialect<affine::AffineDialect,
+                         PimDialect,
                          tensor::TensorDialect,
                          arith::ArithDialect,
                          bufferization::BufferizationDialect,
@@ -217,12 +220,31 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
 
   RewritePatternSet coreBodyPatterns(ctx);
   populateWithGenerated(coreBodyPatterns);
+  populateAffineToStdConversionPatterns(coreBodyPatterns);
   FrozenRewritePatternSet frozenCoreBodyPatterns(std::move(coreBodyPatterns));
 
+  ConversionTarget coreBodyTarget(*ctx);
+  coreBodyTarget.addLegalDialect<PimDialect,
+                                 tensor::TensorDialect,
+                                 arith::ArithDialect,
+                                 bufferization::BufferizationDialect,
+                                 func::FuncDialect,
+                                 memref::MemRefDialect,
+                                 scf::SCFDialect,
+                                 BuiltinDialect>();
+  coreBodyTarget.addLegalOp<spatial::SpatConcatOp,
+                            spatial::SpatChannelReceiveOp,
+                            spatial::SpatChannelReceiveTensorOp,
+                            spatial::SpatChannelReceiveTensorBatchOp,
+                            spatial::SpatChannelSendOp,
+                            spatial::SpatChannelSendTensorOp,
+                            spatial::SpatChannelSendTensorBatchOp,
+                            spatial::SpatExtractRowsOp>();
+
   SmallVector<pim::PimCoreOp> coreOps;
   funcOp.walk([&](pim::PimCoreOp coreOp) { coreOps.push_back(coreOp); });
   for (auto coreOp : coreOps) {
-    if (failed(applyFullConversion(coreOp.getOperation(), target, frozenCoreBodyPatterns))) {
+    if (failed(applyFullConversion(coreOp.getOperation(), coreBodyTarget, frozenCoreBodyPatterns))) {
       coreOp.emitOpError("failed to convert nested Spatial ops inside pim.core");
       signalPassFailure();
       return;
@@ -232,7 +254,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   SmallVector<pim::PimCoreBatchOp> coreBatchOps;
   funcOp.walk([&](pim::PimCoreBatchOp coreBatchOp) { coreBatchOps.push_back(coreBatchOp); });
   for (auto coreBatchOp : coreBatchOps) {
-    if (failed(applyFullConversion(coreBatchOp.getOperation(), target, frozenCoreBodyPatterns))) {
+    if (failed(applyFullConversion(coreBatchOp.getOperation(), coreBodyTarget, frozenCoreBodyPatterns))) {
       coreBatchOp.emitOpError("failed to convert nested Spatial ops inside pim.core_batch");
       signalPassFailure();
       return;