fixes

slightly faster codegen

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp

@@ -2180,39 +2180,55 @@ static Value createIm2colRows(const ConvLoweringState& state,
   auto elemType = preparedInput.type.getElementType();
   auto packedRowType = RankedTensorType::get(
     {plan.effectiveMaxParallelPixels * plan.patchSize}, elemType, plan.gemmInputRowsType.getEncoding());
-  auto zeroAttr = DenseElementsAttr::get(packedRowType, rewriter.getZeroAttr(elemType));
-  Value zeroRow = getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, packedRowType);
+  auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType);
+  auto patchRowType = RankedTensorType::get({plan.patchSize}, elemType);
+  bool hasPartialLane = plan.chunkNumPatches % plan.effectiveMaxParallelPixels != 0;
+  SmallVector<Value> im2colInputs {preparedInput.value};
   auto im2colComputeOp = createSpatComputeBatch(
     rewriter,
     loc,
     TypeRange {plan.gemmInputRowsType},
     plan.packedNumRows,
     {},
-    ValueRange {preparedInput.value, zeroRow},
+    im2colInputs,
     [&](detail::SpatComputeBatchBodyArgs args) {
       Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
       Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
       Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
       Value cPack = getOrCreateIndexConstant(rewriter, anchorOp, plan.effectiveMaxParallelPixels);
-      Value cNumPatches = getOrCreateIndexConstant(rewriter, anchorOp, plan.chunkNumPatches);
       Value laneStart = affineMulConst(rewriter, loc, args.lane, plan.effectiveMaxParallelPixels, anchorOp);
-      Value remaining = arith::SubIOp::create(rewriter, loc, cNumPatches, laneStart);
-      Value isPartial = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::ult, remaining, cPack);
-      Value lanePatches = arith::SelectOp::create(rewriter, loc, isPartial, remaining, cPack);
-      auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType);
-      auto patchRowType = RankedTensorType::get({plan.patchSize}, elemType);
-
+      Value lanePatches = cPack;
+      if (hasPartialLane) {
+        Value cNumPatches = getOrCreateIndexConstant(rewriter, anchorOp, plan.chunkNumPatches);
+        Value remaining = arith::SubIOp::create(rewriter, loc, cNumPatches, laneStart);
+        Value isPartial = arith::CmpIOp::create(
+          rewriter, loc, arith::CmpIPredicate::ult, remaining, cPack);
+        lanePatches = arith::SelectOp::create(rewriter, loc, isPartial, remaining, cPack);
+      }
+      Value rowInit = tensor::EmptyOp::create(rewriter, loc, packedRowType.getShape(), elemType);
+      if (hasPartialLane) {
+        auto zeroAttr = cast<TypedAttr>(rewriter.getZeroAttr(elemType));
+        rowInit = linalg::MapOp::create(
+          rewriter, loc, ValueRange {}, rowInit,
+          [&](OpBuilder& builder, Location nestedLoc, ValueRange) {
+            Value zero = arith::ConstantOp::create(builder, nestedLoc, zeroAttr);
+            linalg::YieldOp::create(builder, nestedLoc, zero);
+          }).getResult().front();
+      }
       auto rowLoop = buildNormalizedScfFor(
         rewriter,
         loc,
         c0,
         lanePatches,
         c1,
-        ValueRange {args.inputs[1]},
+        ValueRange {rowInit},
         [&](OpBuilder&, Location nestedLoc, Value copyIndex, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
           Value patchIndex = arith::AddIOp::create(rewriter, nestedLoc, laneStart, copyIndex);
-          Value batchIndex =
-            affineAddFloorDivConst(rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchorOp);
+          Value batchIndex = state.batchSize == 1
+            ? c0
+            : affineAddFloorDivConst(
+                rewriter, nestedLoc, patchIndex, plan.chunkStart,
+                plan.numPatchesPerBatch, anchorOp);
           Value batchPatchIndex =
             affineAddModConst(rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchorOp);
           Value outHeightIndex = affineFloorDivConst(rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchorOp);
@@ -2249,7 +2265,8 @@ static Value createIm2colRows(const ConvLoweringState& state,
         });
       if (failed(rowLoop))
         return failure();
-      publishGraphBatchPhysicalFragment(rewriter, loc, rowLoop->results.front(), args.outputs.front(), args.lane);
+      Value row = rowLoop->results.front();
+      publishGraphBatchPhysicalFragment(rewriter, loc, row, args.outputs.front(), args.lane);
       return success();
     });
 

src/PIM/Conversion/SpatialToPim/Common.cpp

@@ -1,6 +1,7 @@
 #include "mlir/IR/ValueRange.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 
 #include "llvm/ADT/STLExtras.h"
 
@@ -59,6 +60,23 @@ bool hasLaterUserInBlock(mlir::Value value, Operation* operation) {
   return false;
 }
 
+static bool isTensorView(mlir::Value value) {
+  return isa_and_nonnull<tensor::CastOp,
+                         tensor::CollapseShapeOp,
+                         tensor::ExpandShapeOp,
+                         tensor::ExtractSliceOp,
+                         tensor::ReshapeOp>(value.getDefiningOp());
+}
+
+static bool isLoopCarriedOutput(mlir::Value operand, Operation* operation) {
+  auto argument = dyn_cast<BlockArgument>(operand);
+  if (!argument || argument.getArgNumber() == 0 || operation->getBlock() != argument.getOwner())
+    return false;
+  auto loop = dyn_cast_or_null<scf::ForOp>(argument.getOwner()->getParentOp());
+  return loop && cast<scf::YieldOp>(loop.getBody()->getTerminator())
+                   .getOperand(argument.getArgNumber() - 1) == operation->getResult(0);
+}
+
 mlir::Value getBestOutputTensorFromOperandsOrAllocate(RewriterBase& rewriter, Operation* operation) {
   assert("Only support operations with a single result" && operation->getNumResults() == 1);
   mlir::Value result = operation->getResult(0);
@@ -67,7 +85,11 @@ mlir::Value getBestOutputTensorFromOperandsOrAllocate(RewriterBase& rewriter, Op
 
   SmallVector<mlir::Value> operands = getOpOperandsSortedByUses(operation);
   auto validOperands = make_filter_range(operands, [operation, resultType](mlir::Value operand) {
-    return operand.getType() == resultType && !hasLaterUserInBlock(operand, operation);
+    return operand.getType() == resultType
+      && (!isa<BlockArgument>(operand) || isLoopCarriedOutput(operand, operation))
+      && !operand.getDefiningOp<arith::ConstantOp>()
+      && !isTensorView(operand)
+      && !hasLaterUserInBlock(operand, operation);
   });
   auto bestOperand = validOperands.begin();
 

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -2,6 +2,7 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/Utils/Utils.h"
@@ -112,7 +113,9 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                          memref::MemRefDialect,
                          scf::SCFDialect,
                          BuiltinDialect>();
-  target.addLegalOp<spatial::SpatConcatOp,
+  target.addLegalOp<linalg::MapOp,
+                    linalg::YieldOp,
+                    spatial::SpatConcatOp,
                     spatial::SpatChannelReceiveOp,
                     spatial::SpatChannelSendOp,
                     spatial::SpatExtractRowsOp>();
@@ -186,7 +189,9 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                                  memref::MemRefDialect,
                                  scf::SCFDialect,
                                  BuiltinDialect>();
-  coreBodyTarget.addLegalOp<spatial::SpatConcatOp,
+  coreBodyTarget.addLegalOp<linalg::MapOp,
+                            linalg::YieldOp,
+                            spatial::SpatConcatOp,
                             spatial::SpatChannelReceiveOp,
                             spatial::SpatChannelSendOp,
                             spatial::SpatExtractRowsOp>();
@@ -234,7 +239,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                                       memref::MemRefDialect,
                                       scf::SCFDialect,
                                       BuiltinDialect>();
-  communicationTarget.addLegalOp<ModuleOp>();
+  communicationTarget.addLegalOp<ModuleOp, linalg::MapOp, linalg::YieldOp>();
   communicationTarget.addIllegalOp<spatial::SpatConcatOp,
                                    spatial::SpatChannelReceiveOp,
                                    spatial::SpatChannelSendOp,