add shared loop creation helpers

add shared checked arithmetic helpers
refactor pim passes into Pim/Transforms
more robust memory coalescing pass

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

@@ -15,6 +15,7 @@
 
 #include "Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
@@ -247,16 +248,16 @@ static Value createPaddedInputCompute(Value input,
   return computeOp.getResult(0);
 }
 
-static spatial::SpatComputeBatch createVmmBatch(Value a,
-                                                Value b,
-                                                RankedTensorType aType,
-                                                RankedTensorType paddedBType,
-                                                RankedTensorType partialPiecesType,
-                                                int64_t numOutRows,
-                                                int64_t numKSlices,
-                                                int64_t numOutHSlices,
-                                                ConversionPatternRewriter& rewriter,
-                                                Location loc) {
+static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
+                                                           Value b,
+                                                           RankedTensorType aType,
+                                                           RankedTensorType paddedBType,
+                                                           RankedTensorType partialPiecesType,
+                                                           int64_t numOutRows,
+                                                           int64_t numKSlices,
+                                                           int64_t numOutHSlices,
+                                                           ConversionPatternRewriter& rewriter,
+                                                           Location loc) {
   const int64_t laneCount = partialPiecesType.getDimSize(0);
   auto batchOp = createSpatComputeBatch(
     rewriter,
@@ -294,7 +295,8 @@ static spatial::SpatComputeBatch createVmmBatch(Value a,
       createParallelInsertSliceIntoBatchOutput(
         rewriter, loc, piece, args.outputs.front(), pieceOffsets, pieceSizes, unitStrides);
     });
-  assert(succeeded(batchOp) && "expected Gemm VMM batch construction to succeed");
+  if (failed(batchOp))
+    return failure();
   return *batchOp;
 }
 
@@ -416,15 +418,15 @@ static Value createBroadcastedBiasScalar(Value bias,
   return tensor::SplatOp::create(rewriter, loc, scalarType, scalar).getResult();
 }
 
-static spatial::SpatComputeBatch createVvdmulBatch(Value a,
-                                                   Value b,
-                                                   RankedTensorType aType,
-                                                   RankedTensorType bType,
-                                                   RankedTensorType scalarPiecesType,
-                                                   RankedTensorType outType,
-                                                   bool bAlreadyTransposed,
-                                                   ConversionPatternRewriter& rewriter,
-                                                   Location loc) {
+static FailureOr<spatial::SpatComputeBatch> createVvdmulBatch(Value a,
+                                                              Value b,
+                                                              RankedTensorType aType,
+                                                              RankedTensorType bType,
+                                                              RankedTensorType scalarPiecesType,
+                                                              RankedTensorType outType,
+                                                              bool bAlreadyTransposed,
+                                                              ConversionPatternRewriter& rewriter,
+                                                              Location loc) {
   const int64_t numOutRows = outType.getDimSize(0);
   const int64_t numOutCols = outType.getDimSize(1);
   const int64_t reductionSize = aType.getDimSize(1);
@@ -454,26 +456,27 @@ static spatial::SpatComputeBatch createVvdmulBatch(Value a,
       createParallelInsertSliceIntoBatchOutput(
         rewriter, loc, scalar, args.outputs.front(), outputOffsets, scalarSizes, unitStrides);
     });
-  assert(succeeded(batchOp) && "expected Gemm VVDMul batch construction to succeed");
+  if (failed(batchOp))
+    return failure();
   return *batchOp;
 }
 
-static spatial::SpatCompute createDynamicGemmOutputCompute(Value scalarPieces,
-                                                           Value bias,
-                                                           RankedTensorType scalarPiecesType,
-                                                           RankedTensorType biasType,
-                                                           RankedTensorType outType,
-                                                           float alpha,
-                                                           float beta,
-                                                           ConversionPatternRewriter& rewriter,
-                                                           Location loc) {
+static FailureOr<spatial::SpatCompute> createDynamicGemmOutputCompute(Value scalarPieces,
+                                                                      Value bias,
+                                                                      RankedTensorType scalarPiecesType,
+                                                                      RankedTensorType biasType,
+                                                                      RankedTensorType outType,
+                                                                      float alpha,
+                                                                      float beta,
+                                                                      ConversionPatternRewriter& rewriter,
+                                                                      Location loc) {
   const int64_t laneCount = scalarPiecesType.getDimSize(0);
   const int64_t numOutCols = outType.getDimSize(1);
   SmallVector<Value> inputs {scalarPieces};
   if (bias)
     inputs.push_back(bias);
 
-  return createSpatCompute(rewriter, loc, TypeRange {outType}, {}, inputs, [&](ValueRange blockArgs) {
+  return createSpatCompute(rewriter, loc, TypeRange {outType}, {}, inputs, [&](ValueRange blockArgs) -> LogicalResult {
     Value pieces = blockArgs[0];
     Value biasArg = bias ? blockArgs[1] : Value();
     auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
@@ -481,40 +484,50 @@ static spatial::SpatCompute createDynamicGemmOutputCompute(Value scalarPieces,
     Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
     Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1);
     Value cLaneCount = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), laneCount);
-    auto loop = scf::ForOp::create(rewriter, loc, c0, cLaneCount, c1, ValueRange {outputInit});
-    rewriter.setInsertionPointToStart(loop.getBody());
+    auto loop = buildNormalizedScfFor(
+      rewriter,
+      loc,
+      c0,
+      cLaneCount,
+      c1,
+      ValueRange {outputInit},
+      [&](OpBuilder&, Location nestedLoc, Value lane, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+        Value outputAcc = iterArgs.front();
+        Value row = createDynamicGemmBatchRow(lane, numOutCols, rewriter, nestedLoc);
+        Value column =
+          onnx_mlir::affineModConst(rewriter, nestedLoc, lane, numOutCols, rewriter.getInsertionBlock()->getParentOp());
+        SmallVector<OpFoldResult> scalarOffsets {lane, rewriter.getIndexAttr(0)};
+        SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+        SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+        Value scalar = tensor::ExtractSliceOp::create(
+                         rewriter, nestedLoc, scalarType, pieces, scalarOffsets, scalarSizes, unitStrides)
+                         .getResult();
+        if (alpha != 1.0f) {
+          Value alphaTensor = createScalarTensorConstant(scalarType, alpha, rewriter, nestedLoc);
+          scalar = spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, scalar, alphaTensor).getResult();
+        }
+        if (biasArg) {
+          Value biasScalar =
+            createBroadcastedBiasScalar(biasArg, biasType, row, column, scalarType, rewriter, nestedLoc);
+          if (beta != 1.0f) {
+            Value betaTensor = createScalarTensorConstant(scalarType, beta, rewriter, nestedLoc);
+            biasScalar =
+              spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, biasScalar, betaTensor).getResult();
+          }
+          scalar = spatial::SpatVAddOp::create(rewriter, nestedLoc, scalarType, scalar, biasScalar).getResult();
+        }
+        SmallVector<OpFoldResult> outputOffsets {row, column};
+        Value outputNext =
+          tensor::InsertSliceOp::create(rewriter, nestedLoc, scalar, outputAcc, outputOffsets, scalarSizes, unitStrides)
+            .getResult();
+        yielded.push_back(outputNext);
+        return success();
+      });
+    if (failed(loop))
+      return failure();
 
-    Value lane = loop.getInductionVar();
-    Value outputAcc = loop.getRegionIterArgs().front();
-    Value row = createDynamicGemmBatchRow(lane, numOutCols, rewriter, loc);
-    Value column =
-      onnx_mlir::affineModConst(rewriter, loc, lane, numOutCols, rewriter.getInsertionBlock()->getParentOp());
-    SmallVector<OpFoldResult> scalarOffsets {lane, rewriter.getIndexAttr(0)};
-    SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-    SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-    Value scalar =
-      tensor::ExtractSliceOp::create(rewriter, loc, scalarType, pieces, scalarOffsets, scalarSizes, unitStrides)
-        .getResult();
-    if (alpha != 1.0f) {
-      Value alphaTensor = createScalarTensorConstant(scalarType, alpha, rewriter, loc);
-      scalar = spatial::SpatVMulOp::create(rewriter, loc, scalarType, scalar, alphaTensor).getResult();
-    }
-    if (biasArg) {
-      Value biasScalar = createBroadcastedBiasScalar(biasArg, biasType, row, column, scalarType, rewriter, loc);
-      if (beta != 1.0f) {
-        Value betaTensor = createScalarTensorConstant(scalarType, beta, rewriter, loc);
-        biasScalar = spatial::SpatVMulOp::create(rewriter, loc, scalarType, biasScalar, betaTensor).getResult();
-      }
-      scalar = spatial::SpatVAddOp::create(rewriter, loc, scalarType, scalar, biasScalar).getResult();
-    }
-    SmallVector<OpFoldResult> outputOffsets {row, column};
-    Value outputNext =
-      tensor::InsertSliceOp::create(rewriter, loc, scalar, outputAcc, outputOffsets, scalarSizes, unitStrides)
-        .getResult();
-    scf::YieldOp::create(rewriter, loc, outputNext);
-
-    rewriter.setInsertionPointAfter(loop);
-    spatial::SpatYieldOp::create(rewriter, loc, loop.getResult(0));
+    spatial::SpatYieldOp::create(rewriter, loc, loop->results.front());
+    return success();
   });
 }
 
@@ -579,85 +592,92 @@ static Value reducePartialPiecesForHSlice(Value partialPiecesArg,
   return activePieces.front();
 }
 
-static spatial::SpatCompute createReductionCompute(Value partialPieces,
-                                                   Value bias,
-                                                   RankedTensorType partialPiecesType,
-                                                   RankedTensorType outType,
-                                                   RankedTensorType paddedOutType,
-                                                   int64_t numKSlices,
-                                                   ConversionPatternRewriter& rewriter,
-                                                   Location loc) {
+static FailureOr<spatial::SpatCompute> createReductionCompute(Value partialPieces,
+                                                              Value bias,
+                                                              RankedTensorType partialPiecesType,
+                                                              RankedTensorType outType,
+                                                              RankedTensorType paddedOutType,
+                                                              int64_t numKSlices,
+                                                              ConversionPatternRewriter& rewriter,
+                                                              Location loc) {
   SmallVector<Value> inputs {partialPieces};
   if (bias)
     inputs.push_back(bias);
 
-  auto computeOp = createSpatCompute(rewriter, loc, TypeRange {outType}, {}, inputs, [&](ValueRange blockArgs) {
-    Value partialPiecesArg = blockArgs[0];
-    Value biasArg = bias ? blockArgs[1] : Value();
-    if (biasArg && cast<RankedTensorType>(biasArg.getType()) != paddedOutType)
-      biasArg = createZeroPaddedTensor(biasArg, paddedOutType, rewriter, loc);
+  auto computeOp =
+    createSpatCompute(rewriter, loc, TypeRange {outType}, {}, inputs, [&](ValueRange blockArgs) -> LogicalResult {
+      Value partialPiecesArg = blockArgs[0];
+      Value biasArg = bias ? blockArgs[1] : Value();
+      if (biasArg && cast<RankedTensorType>(biasArg.getType()) != paddedOutType)
+        biasArg = createZeroPaddedTensor(biasArg, paddedOutType, rewriter, loc);
 
-    const int64_t numOutRows = outType.getDimSize(0);
-    const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(1), crossbarSize.getValue());
-    auto pieceType = RankedTensorType::get({numOutRows, static_cast<int64_t>(crossbarSize.getValue())},
-                                           partialPiecesType.getElementType());
+      const int64_t numOutRows = outType.getDimSize(0);
+      const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(1), crossbarSize.getValue());
+      auto pieceType = RankedTensorType::get({numOutRows, static_cast<int64_t>(crossbarSize.getValue())},
+                                             partialPiecesType.getElementType());
 
-    Value outputInit =
-      tensor::EmptyOp::create(rewriter, loc, paddedOutType.getShape(), paddedOutType.getElementType()).getResult();
-    SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-    SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(numOutRows),
-                                          rewriter.getIndexAttr(crossbarSize.getValue())};
+      Value outputInit =
+        tensor::EmptyOp::create(rewriter, loc, paddedOutType.getShape(), paddedOutType.getElementType()).getResult();
+      SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+      SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(numOutRows),
+                                            rewriter.getIndexAttr(crossbarSize.getValue())};
 
-    auto buildOutputSlice = [&](Value outputAcc, Value hSlice) -> Value {
-      Value reduced =
-        reducePartialPiecesForHSlice(partialPiecesArg, hSlice, pieceType, numKSlices, numOutRows, rewriter, loc);
-      Value hOffset = onnx_mlir::affineMulConst(
-        rewriter, loc, hSlice, crossbarSize.getValue(), rewriter.getInsertionBlock()->getParentOp());
-      if (biasArg) {
-        SmallVector<OpFoldResult> biasOffsets {rewriter.getIndexAttr(0), hOffset};
-        Value biasSlice =
-          tensor::ExtractSliceOp::create(rewriter, loc, pieceType, biasArg, biasOffsets, pieceSizes, unitStrides)
-            .getResult();
-        reduced = spatial::SpatVAddOp::create(rewriter, loc, pieceType, reduced, biasSlice).getResult();
+      auto buildOutputSlice = [&](Value outputAcc, Value hSlice) -> Value {
+        Value reduced =
+          reducePartialPiecesForHSlice(partialPiecesArg, hSlice, pieceType, numKSlices, numOutRows, rewriter, loc);
+        Value hOffset = onnx_mlir::affineMulConst(
+          rewriter, loc, hSlice, crossbarSize.getValue(), rewriter.getInsertionBlock()->getParentOp());
+        if (biasArg) {
+          SmallVector<OpFoldResult> biasOffsets {rewriter.getIndexAttr(0), hOffset};
+          Value biasSlice =
+            tensor::ExtractSliceOp::create(rewriter, loc, pieceType, biasArg, biasOffsets, pieceSizes, unitStrides)
+              .getResult();
+          reduced = spatial::SpatVAddOp::create(rewriter, loc, pieceType, reduced, biasSlice).getResult();
+        }
+
+        SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), hOffset};
+        return tensor::InsertSliceOp::create(rewriter, loc, reduced, outputAcc, outputOffsets, pieceSizes, unitStrides)
+          .getResult();
+      };
+
+      Value paddedOutput = outputInit;
+      if (numOutHSlices == 1) {
+        Value hSlice = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
+        paddedOutput = buildOutputSlice(outputInit, hSlice);
+      }
+      else {
+        Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
+        Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1);
+        Value cOutHSlices =
+          getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutHSlices);
+        auto hLoop = buildNormalizedScfFor(
+          rewriter,
+          loc,
+          c0,
+          cOutHSlices,
+          c1,
+          ValueRange {outputInit},
+          [&](OpBuilder&, Location, Value hSlice, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+            yielded.push_back(buildOutputSlice(iterArgs.front(), hSlice));
+            return success();
+          });
+        if (failed(hLoop))
+          return failure();
+        paddedOutput = hLoop->results.front();
       }
 
-      SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), hOffset};
-      return tensor::InsertSliceOp::create(rewriter, loc, reduced, outputAcc, outputOffsets, pieceSizes, unitStrides)
-        .getResult();
-    };
-
-    Value paddedOutput = outputInit;
-    if (numOutHSlices == 1) {
-      Value hSlice = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
-      paddedOutput = buildOutputSlice(outputInit, hSlice);
-    }
-    else {
-      Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
-      Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1);
-      Value cOutHSlices =
-        getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutHSlices);
-      auto hLoop = scf::ForOp::create(rewriter, loc, c0, cOutHSlices, c1, ValueRange {outputInit});
-      rewriter.setInsertionPointToStart(hLoop.getBody());
-
-      Value hSlice = hLoop.getInductionVar();
-      Value outputAcc = hLoop.getRegionIterArgs().front();
-      scf::YieldOp::create(rewriter, loc, buildOutputSlice(outputAcc, hSlice));
-
-      rewriter.setInsertionPointAfter(hLoop);
-      paddedOutput = hLoop.getResult(0);
-    }
-
-    Value result = paddedOutput;
-    if (paddedOutType != outType) {
-      SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(outType.getDimSize(0)),
-                                             rewriter.getIndexAttr(outType.getDimSize(1))};
-      result =
-        tensor::ExtractSliceOp::create(rewriter, loc, outType, paddedOutput, outputOffsets, outputSizes, unitStrides)
-          .getResult();
-    }
-    spatial::SpatYieldOp::create(rewriter, loc, result);
-  });
+      Value result = paddedOutput;
+      if (paddedOutType != outType) {
+        SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+        SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(outType.getDimSize(0)),
+                                               rewriter.getIndexAttr(outType.getDimSize(1))};
+        result =
+          tensor::ExtractSliceOp::create(rewriter, loc, outType, paddedOutput, outputOffsets, outputSizes, unitStrides)
+            .getResult();
+      }
+      spatial::SpatYieldOp::create(rewriter, loc, result);
+      return success();
+    });
 
   return computeOp;
 }
@@ -755,9 +775,13 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
     auto scalarPiecesType = RankedTensorType::get({laneCount64, 1}, outType.getElementType());
     auto batchOp =
       createVvdmulBatch(a, b, aType, bType, scalarPiecesType, outType, gemmOpAdaptor.getTransB(), rewriter, loc);
+    if (failed(batchOp))
+      return failure();
     auto outputCompute = createDynamicGemmOutputCompute(
-      batchOp.getResult(0), hasC ? c : Value(), scalarPiecesType, biasType, outType, alpha, beta, rewriter, loc);
-    rewriter.replaceOp(gemmOp, outputCompute.getResults());
+      batchOp->getResult(0), hasC ? c : Value(), scalarPiecesType, biasType, outType, alpha, beta, rewriter, loc);
+    if (failed(outputCompute))
+      return failure();
+    rewriter.replaceOp(gemmOp, outputCompute->getResults());
     return success();
   }
 
@@ -832,10 +856,14 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
     RankedTensorType::get({laneCount64, static_cast<int64_t>(crossbarSize.getValue())}, outType.getElementType());
   auto batchOp =
     createVmmBatch(a, b, aType, paddedBType, partialPiecesType, numOutRows, numKSlices, numOutHSlices, rewriter, loc);
+  if (failed(batchOp))
+    return failure();
   auto reductionCompute = createReductionCompute(
-    batchOp.getResult(0), bias, partialPiecesType, outType, paddedOutType, numKSlices, rewriter, loc);
+    batchOp->getResult(0), bias, partialPiecesType, outType, paddedOutType, numKSlices, rewriter, loc);
+  if (failed(reductionCompute))
+    return failure();
 
-  rewriter.replaceOp(gemmOp, reductionCompute.getResults());
+  rewriter.replaceOp(gemmOp, reductionCompute->getResults());
   return success();
 }