finish helper refactoring

use uniqued constant helpers everywhere
materialize transposed constants directly

src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp

@@ -7,6 +7,7 @@
 #include "mlir/IR/Matchers.h"
 
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
 #include "Conversion/SpatialToPim/SpatialToPimPass.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
@@ -18,15 +19,9 @@ using namespace onnx_mlir::pim;
 namespace onnx_mlir {
 namespace {
 
-static bool isExplicitHostOperand(Operation* op, unsigned operandIndex) {
-  if (isa<pim::PimMemCopyDevToHostOp>(op))
-    return operandIndex == 2;
-  return false;
-}
-
 static bool isUsedOnlyAsExplicitHostOperand(Value value) {
   return !value.use_empty() && llvm::all_of(value.getUses(), [](OpOperand& use) {
-    return isExplicitHostOperand(use.getOwner(), use.getOperandNumber());
+    return isExplicitDevToHostTargetOperand(use.getOwner(), use.getOperandNumber());
   });
 }
 
@@ -55,7 +50,7 @@ static Value createScaledIndexValue(IRRewriter& rewriter, Location loc, Value ba
   if (scale == 1)
     return base;
 
-  auto scaleValue = arith::ConstantIndexOp::create(rewriter, loc, scale).getResult();
+  auto scaleValue = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), scale);
   return arith::MulIOp::create(rewriter, loc, base, scaleValue).getResult();
 }
 
@@ -77,7 +72,8 @@ static Value createHostTargetOffset(IRRewriter& rewriter,
     if (auto attr = dyn_cast<Attribute>(offset)) {
       auto intAttr = dyn_cast<IntegerAttr>(attr);
       assert(intAttr && "expected integer offset attribute");
-      scaledOffset = arith::ConstantIndexOp::create(rewriter, loc, intAttr.getInt() * scale).getResult();
+      scaledOffset =
+        getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), intAttr.getInt() * scale);
     }
     else {
       scaledOffset = createScaledIndexValue(rewriter, loc, mapper.lookupOrDefault(cast<Value>(offset)), scale);
@@ -88,7 +84,7 @@ static Value createHostTargetOffset(IRRewriter& rewriter,
   }
 
   if (!totalOffset)
-    totalOffset = arith::ConstantIndexOp::create(rewriter, loc, 0).getResult();
+    totalOffset = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
   return totalOffset;
 }
 
@@ -214,7 +210,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
         Value hostTarget = getOrCreateHostOutputTensor(resultIndex, insertSlice.getLoc());
         auto hostTargetType = cast<ShapedType>(hostTarget.getType());
         Value hostTargetOffset = createHostTargetOffset(rewriter, insertSlice, hostTargetType, mapper);
-        Value zeroOffset = arith::ConstantIndexOp::create(rewriter, insertSlice.getLoc(), 0).getResult();
+        Value zeroOffset = getOrCreateIndexConstant(rewriter, coreBatchOp.getOperation(), 0);
         pim::PimMemCopyDevToHostOp::create(rewriter,
                                            insertSlice.getLoc(),
                                            hostTarget.getType(),
@@ -254,7 +250,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
     for (auto [operandIndex, operand] : llvm::enumerate(op.getOperands())) {
       if (!isa<TensorType>(operand.getType()) || mapper.contains(operand))
         continue;
-      if (isExplicitHostOperand(&op, operandIndex))
+      if (isExplicitDevToHostTargetOperand(&op, operandIndex))
         continue;
 
       Operation* definingOp = operand.getDefiningOp();

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -40,7 +40,7 @@ cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewriter& rewrite
       continue;
 
     if (auto constantOp = dyn_cast<arith::ConstantOp>(definingOp)) {
-      mapping.map(operand, getOrCreateHostConstantLike(constantFolder, constantOp));
+      mapping.map(operand, getOrCreateConstantLike(constantFolder, constantOp));
       continue;
     }
 
@@ -218,7 +218,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatCompute comp
       continue;
 
     if (auto constantOp = input.getDefiningOp<arith::ConstantOp>()) {
-      blockArg->replaceAllUsesWith(getOrCreateHostConstantLike(constantFolder, constantOp));
+      blockArg->replaceAllUsesWith(getOrCreateConstantLike(constantFolder, constantOp));
       continue;
     }
 
@@ -230,8 +230,8 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatCompute comp
       PimMemCopyHostToDevOp::create(rewriter,
                                     loc,
                                     outputBuffer.getType(),
-                                    getOrCreateHostIndexConstant(constantFolder, outputBuffer.getOperation(), 0),
-                                    getOrCreateHostIndexConstant(constantFolder, outputBuffer.getOperation(), 0),
+                                    getOrCreateIndexConstant(constantFolder, outputBuffer.getOperation(), 0),
+                                    getOrCreateIndexConstant(constantFolder, outputBuffer.getOperation(), 0),
                                     outputBuffer,
                                     input,
                                     getTensorSizeInBytesAttr(rewriter, input))

src/PIM/Conversion/SpatialToPim/Patterns.cpp

@@ -16,25 +16,9 @@ void populateInitialPatterns(RewritePatternSet& patterns) {
   populateTransposeLoweringPatterns(patterns);
 }
 
-void populateGlobalTensorMaterializationPatternPhase(RewritePatternSet& patterns) {
-  populateGlobalTensorMaterializationPatterns(patterns);
-}
-
-void populateInitialTensorPackingPatterns(RewritePatternSet& patterns) {
-  populateTensorPackingPatterns(patterns);
-}
-
 void populateCoreBodyPatterns(RewritePatternSet& patterns) {
   raptor::populateWithGenerated(patterns);
   populateTransposeLoweringPatterns(patterns);
 }
 
-void populateFinalTensorPackingPatterns(RewritePatternSet& patterns) {
-  populateTensorPackingPatterns(patterns);
-}
-
-void populateCommunicationPatterns(RewritePatternSet& patterns) {
-  populateChannelLoweringPatterns(patterns);
-}
-
 } // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/Patterns.hpp

@@ -9,11 +9,7 @@
 namespace onnx_mlir {
 
 void populateInitialPatterns(mlir::RewritePatternSet& patterns);
-void populateGlobalTensorMaterializationPatternPhase(mlir::RewritePatternSet& patterns);
-void populateInitialTensorPackingPatterns(mlir::RewritePatternSet& patterns);
 void populateCoreBodyPatterns(mlir::RewritePatternSet& patterns);
-void populateFinalTensorPackingPatterns(mlir::RewritePatternSet& patterns);
-void populateCommunicationPatterns(mlir::RewritePatternSet& patterns);
 
 void populateTransposeLoweringPatterns(mlir::RewritePatternSet& patterns);
 void populateChannelLoweringPatterns(mlir::RewritePatternSet& patterns);

src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp

@@ -326,7 +326,7 @@ cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewriter& rewrite
       continue;
 
     if (auto constantOp = dyn_cast<arith::ConstantOp>(definingOp)) {
-      mapping.map(operand, getOrCreateHostConstantLike(constantFolder, constantOp));
+      mapping.map(operand, getOrCreateConstantLike(constantFolder, constantOp));
       continue;
     }
 
@@ -370,8 +370,8 @@ static Value emitHostCopy(IRRewriter& rewriter,
                           OperationFolder& constantFolder) {
   Operation* anchorOp = sourceValue.getDefiningOp() ? sourceValue.getDefiningOp() : outputTensor.getDefiningOp();
   assert(anchorOp && "expected a concrete op anchor for return-path host copy constants");
-  Value hostTargetOffsetValue = getOrCreateHostIndexConstant(constantFolder, anchorOp, hostTargetOffset);
-  Value deviceSourceOffsetValue = getOrCreateHostIndexConstant(constantFolder, anchorOp, deviceSourceOffset);
+  Value hostTargetOffsetValue = getOrCreateIndexConstant(constantFolder, anchorOp, hostTargetOffset);
+  Value deviceSourceOffsetValue = getOrCreateIndexConstant(constantFolder, anchorOp, deviceSourceOffset);
   return PimMemCopyDevToHostOp::create(rewriter,
                                        loc,
                                        outputTensor.getType(),

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -81,7 +81,7 @@ static Value createZeroedDeviceHVector(IRRewriter& rewriter,
   auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, tensorType);
   auto zeroGlobal = getOrCreateZeroGlobal(rewriter, loc, tensorType);
   auto zeroValue = memref::GetGlobalOp::create(rewriter, loc, zeroGlobal.getType(), zeroGlobal.getName());
-  auto zeroIndex = getOrCreateHostIndexConstant(constantFolder, outputBuffer.getOperation(), 0);
+  auto zeroIndex = getOrCreateIndexConstant(constantFolder, outputBuffer.getOperation(), 0);
   auto sizeAttr = rewriter.getI32IntegerAttr(static_cast<int32_t>(getShapedTypeSizeInBytes(tensorType)));
 
   if (outputBuffer->getParentOfType<PimCoreBatchOp>())
@@ -160,7 +160,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   }
 
   RewritePatternSet globalTensorPatterns(ctx);
-  populateGlobalTensorMaterializationPatternPhase(globalTensorPatterns);
+  populateGlobalTensorMaterializationPatterns(globalTensorPatterns);
   walkAndApplyPatterns(moduleOp, std::move(globalTensorPatterns));
 
   auto returnOp = cast<func::ReturnOp>(funcOp.front().getTerminator());
@@ -190,7 +190,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   }
 
   RewritePatternSet initialTensorPackingPatterns(ctx);
-  populateInitialTensorPackingPatterns(initialTensorPackingPatterns);
+  populateTensorPackingPatterns(initialTensorPackingPatterns);
   walkAndApplyPatterns(funcOp, std::move(initialTensorPackingPatterns));
   eraseUnusedTensorPackingOps(funcOp, rewriter);
 
@@ -250,7 +250,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   eraseOpsToRemove();
 
   RewritePatternSet finalTensorPackingPatterns(ctx);
-  populateFinalTensorPackingPatterns(finalTensorPackingPatterns);
+  populateTensorPackingPatterns(finalTensorPackingPatterns);
   walkAndApplyPatterns(funcOp, std::move(finalTensorPackingPatterns));
   eraseUnusedTensorPackingOps(funcOp, rewriter);
 
@@ -270,7 +270,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                                    spatial::SpatExtractRowsOp>();
 
   RewritePatternSet communicationPatterns(ctx);
-  populateCommunicationPatterns(communicationPatterns);
+  populateChannelLoweringPatterns(communicationPatterns);
   if (failed(applyFullConversion(funcOp, communicationTarget, std::move(communicationPatterns)))) {
     funcOp.emitOpError("failed to lower Spatial communication ops to PIM communication ops");
     signalPassFailure();
@@ -333,8 +333,8 @@ LogicalResult raptor::SpatialToPimPass::allocateAndInitializeCoreLocalVariables(
       rewriter,
       loc,
       tensorType,
-      getOrCreateHostIndexConstant(constantFolder, deviceTensor.getOperation(), 0),
-      getOrCreateHostIndexConstant(constantFolder,
+      getOrCreateIndexConstant(constantFolder, deviceTensor.getOperation(), 0),
+      getOrCreateIndexConstant(constantFolder,
         deviceTensor.getOperation(), static_cast<int64_t>(elementsOffset * elementByteSize) ),
       deviceTensor,
       inputTensor,