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/NN/Pool.cpp

@@ -12,6 +12,7 @@
 #include <optional>
 #include <type_traits>
 
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
@@ -275,86 +276,102 @@ struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
         Value cStrideHeight = getOrCreateIndexConstant(rewriter, anchorOp, strideHeight);
         Value cStrideWidth = getOrCreateIndexConstant(rewriter, anchorOp, strideWidth);
 
-        auto outputLoop = scf::ForOp::create(rewriter, loc, c0, cOutputPatchCount, c1, ValueRange {pooledOutputInit});
-        rewriter.setInsertionPointToStart(outputLoop.getBody());
+        auto outputLoop = buildNormalizedScfFor(
+          rewriter,
+          loc,
+          c0,
+          cOutputPatchCount,
+          c1,
+          ValueRange {pooledOutputInit},
+          [&](OpBuilder&,
+              Location nestedLoc,
+              Value outputPatchIndex,
+              ValueRange iterArgs,
+              SmallVectorImpl<Value>& yielded) {
+            Value pooledOutputAcc = iterArgs.front();
+            Value batchIndex = arith::DivUIOp::create(rewriter, nestedLoc, outputPatchIndex, cOutputPixelsPerBatch);
+            Value batchPatchIndex =
+              arith::RemUIOp::create(rewriter, nestedLoc, outputPatchIndex, cOutputPixelsPerBatch);
+            Value outHeightIndex = arith::DivUIOp::create(rewriter, nestedLoc, batchPatchIndex, cOutputWidth);
+            Value outWidthIndex = arith::RemUIOp::create(rewriter, nestedLoc, batchPatchIndex, cOutputWidth);
+            Value windowBaseH = arith::MulIOp::create(rewriter, nestedLoc, outHeightIndex, cStrideHeight);
+            Value windowBaseW = arith::MulIOp::create(rewriter, nestedLoc, outWidthIndex, cStrideWidth);
 
-        Value outputPatchIndex = outputLoop.getInductionVar();
-        Value pooledOutputAcc = outputLoop.getRegionIterArgs().front();
+            Value updatedOutput = pooledOutputAcc;
+            for (int64_t channelTile = 0; channelTile < channelTileCount; ++channelTile) {
+              const int64_t tileChannels = std::min<int64_t>(xbarSize, channels - channelTile * xbarSize);
+              auto tileType = RankedTensorType::get({1, tileChannels, 1, 1}, outType.getElementType());
+              Value reducedWindow =
+                createPoolFillTensor(rewriter, nestedLoc, tileType, std::is_same_v<PoolOp, ONNXMaxPoolSingleOutOp>);
 
-        Value batchIndex = arith::DivUIOp::create(rewriter, loc, outputPatchIndex, cOutputPixelsPerBatch);
-        Value batchPatchIndex = arith::RemUIOp::create(rewriter, loc, outputPatchIndex, cOutputPixelsPerBatch);
-        Value outHeightIndex = arith::DivUIOp::create(rewriter, loc, batchPatchIndex, cOutputWidth);
-        Value outWidthIndex = arith::RemUIOp::create(rewriter, loc, batchPatchIndex, cOutputWidth);
-        Value windowBaseH = arith::MulIOp::create(rewriter, loc, outHeightIndex, cStrideHeight);
-        Value windowBaseW = arith::MulIOp::create(rewriter, loc, outWidthIndex, cStrideWidth);
+              for (int64_t kernelH = 0; kernelH < kernelHeight; ++kernelH) {
+                Value paddedInH = windowBaseH;
+                if (kernelH * dilationHeight != 0) {
+                  Value kernelHOffset = getOrCreateIndexConstant(rewriter, anchorOp, kernelH * dilationHeight);
+                  paddedInH = arith::AddIOp::create(rewriter, nestedLoc, paddedInH, kernelHOffset);
+                }
 
-        Value updatedOutput = pooledOutputAcc;
-        for (int64_t channelTile = 0; channelTile < channelTileCount; ++channelTile) {
-          const int64_t tileChannels = std::min<int64_t>(xbarSize, channels - channelTile * xbarSize);
-          auto tileType = RankedTensorType::get({1, tileChannels, 1, 1}, outType.getElementType());
-          Value reducedWindow =
-            createPoolFillTensor(rewriter, loc, tileType, std::is_same_v<PoolOp, ONNXMaxPoolSingleOutOp>);
+                for (int64_t kernelW = 0; kernelW < kernelWidth; ++kernelW) {
+                  Value paddedInW = windowBaseW;
+                  if (kernelW * dilationWidth != 0) {
+                    Value kernelWOffset = getOrCreateIndexConstant(rewriter, anchorOp, kernelW * dilationWidth);
+                    paddedInW = arith::AddIOp::create(rewriter, nestedLoc, paddedInW, kernelWOffset);
+                  }
 
-          for (int64_t kernelH = 0; kernelH < kernelHeight; ++kernelH) {
-            Value paddedInH = windowBaseH;
-            if (kernelH * dilationHeight != 0) {
-              Value kernelHOffset = getOrCreateIndexConstant(rewriter, anchorOp, kernelH * dilationHeight);
-              paddedInH = arith::AddIOp::create(rewriter, loc, paddedInH, kernelHOffset);
-            }
-
-            for (int64_t kernelW = 0; kernelW < kernelWidth; ++kernelW) {
-              Value paddedInW = windowBaseW;
-              if (kernelW * dilationWidth != 0) {
-                Value kernelWOffset = getOrCreateIndexConstant(rewriter, anchorOp, kernelW * dilationWidth);
-                paddedInW = arith::AddIOp::create(rewriter, loc, paddedInW, kernelWOffset);
+                  SmallVector<OpFoldResult> offsets = {
+                    batchIndex, rewriter.getIndexAttr(channelTile * xbarSize), paddedInH, paddedInW};
+                  SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(1),
+                                                     rewriter.getIndexAttr(tileChannels),
+                                                     rewriter.getIndexAttr(1),
+                                                     rewriter.getIndexAttr(1)};
+                  SmallVector<OpFoldResult> strides = {rewriter.getIndexAttr(1),
+                                                       rewriter.getIndexAttr(1),
+                                                       rewriter.getIndexAttr(1),
+                                                       rewriter.getIndexAttr(1)};
+                  Value windowValue =
+                    tensor::ExtractSliceOp::create(rewriter, nestedLoc, tileType, paddedInput, offsets, sizes, strides);
+                  windowValue = materializeTileTensor(rewriter, nestedLoc, windowValue);
+                  reducedWindow = ReduceOp::create(rewriter, nestedLoc, tileType, reducedWindow, windowValue);
+                }
               }
 
-              SmallVector<OpFoldResult> offsets = {
-                batchIndex, rewriter.getIndexAttr(channelTile * xbarSize), paddedInH, paddedInW};
-              SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(1),
-                                                 rewriter.getIndexAttr(tileChannels),
-                                                 rewriter.getIndexAttr(1),
-                                                 rewriter.getIndexAttr(1)};
-              SmallVector<OpFoldResult> strides = {
+              if constexpr (std::is_same_v<PoolOp, ONNXAveragePoolOp>) {
+                SmallVector<OpFoldResult> scaleOffsets = {rewriter.getIndexAttr(0),
+                                                          rewriter.getIndexAttr(channelTile * xbarSize),
+                                                          outHeightIndex,
+                                                          outWidthIndex};
+                SmallVector<OpFoldResult> scaleSizes = {rewriter.getIndexAttr(1),
+                                                        rewriter.getIndexAttr(tileChannels),
+                                                        rewriter.getIndexAttr(1),
+                                                        rewriter.getIndexAttr(1)};
+                SmallVector<OpFoldResult> scaleStrides = {rewriter.getIndexAttr(1),
+                                                          rewriter.getIndexAttr(1),
+                                                          rewriter.getIndexAttr(1),
+                                                          rewriter.getIndexAttr(1)};
+                Value scaleSlice = tensor::ExtractSliceOp::create(
+                  rewriter, nestedLoc, tileType, averageScaleTensor, scaleOffsets, scaleSizes, scaleStrides);
+                scaleSlice = materializeTileTensor(rewriter, nestedLoc, scaleSlice);
+                reducedWindow = spatial::SpatVMulOp::create(rewriter, nestedLoc, tileType, reducedWindow, scaleSlice);
+              }
+
+              SmallVector<OpFoldResult> outputOffsets = {
+                batchIndex, rewriter.getIndexAttr(channelTile * xbarSize), outHeightIndex, outWidthIndex};
+              SmallVector<OpFoldResult> outputSizes = {rewriter.getIndexAttr(1),
+                                                       rewriter.getIndexAttr(tileChannels),
+                                                       rewriter.getIndexAttr(1),
+                                                       rewriter.getIndexAttr(1)};
+              SmallVector<OpFoldResult> outputStrides = {
                 rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-              Value windowValue =
-                tensor::ExtractSliceOp::create(rewriter, loc, tileType, paddedInput, offsets, sizes, strides);
-              windowValue = materializeTileTensor(rewriter, loc, windowValue);
-              reducedWindow = ReduceOp::create(rewriter, loc, tileType, reducedWindow, windowValue);
+              updatedOutput = tensor::InsertSliceOp::create(
+                rewriter, nestedLoc, reducedWindow, updatedOutput, outputOffsets, outputSizes, outputStrides);
             }
-          }
+            yielded.push_back(updatedOutput);
+            return success();
+          });
+        if (failed(outputLoop))
+          return failure();
 
-          if constexpr (std::is_same_v<PoolOp, ONNXAveragePoolOp>) {
-            SmallVector<OpFoldResult> scaleOffsets = {
-              rewriter.getIndexAttr(0), rewriter.getIndexAttr(channelTile * xbarSize), outHeightIndex, outWidthIndex};
-            SmallVector<OpFoldResult> scaleSizes = {rewriter.getIndexAttr(1),
-                                                    rewriter.getIndexAttr(tileChannels),
-                                                    rewriter.getIndexAttr(1),
-                                                    rewriter.getIndexAttr(1)};
-            SmallVector<OpFoldResult> scaleStrides = {
-              rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-            Value scaleSlice = tensor::ExtractSliceOp::create(
-              rewriter, loc, tileType, averageScaleTensor, scaleOffsets, scaleSizes, scaleStrides);
-            scaleSlice = materializeTileTensor(rewriter, loc, scaleSlice);
-            reducedWindow = spatial::SpatVMulOp::create(rewriter, loc, tileType, reducedWindow, scaleSlice);
-          }
-
-          SmallVector<OpFoldResult> outputOffsets = {
-            batchIndex, rewriter.getIndexAttr(channelTile * xbarSize), outHeightIndex, outWidthIndex};
-          SmallVector<OpFoldResult> outputSizes = {rewriter.getIndexAttr(1),
-                                                   rewriter.getIndexAttr(tileChannels),
-                                                   rewriter.getIndexAttr(1),
-                                                   rewriter.getIndexAttr(1)};
-          SmallVector<OpFoldResult> outputStrides = {
-            rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-          updatedOutput = tensor::InsertSliceOp::create(
-            rewriter, loc, reducedWindow, updatedOutput, outputOffsets, outputSizes, outputStrides);
-        }
-
-        scf::YieldOp::create(rewriter, loc, updatedOutput);
-
-        rewriter.setInsertionPointAfter(outputLoop);
-        spatial::SpatYieldOp::create(rewriter, loc, outputLoop.getResult(0));
+        spatial::SpatYieldOp::create(rewriter, loc, outputLoop->results.front());
         return success();
       });
     if (failed(computeOp))

src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Softmax.cpp

@@ -3,6 +3,7 @@
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/DialectConversion.h"
 
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -42,13 +43,13 @@ static Value buildLoopSoftmaxSlice(Value input,
   return tensor::InsertSliceOp::create(rewriter, loc, softmaxSlice, accumulator, offsets, sizes, strides);
 }
 
-static Value buildLoopSoftmaxNest(Value input,
-                                  Value accumulator,
-                                  RankedTensorType inputType,
-                                  int64_t axis,
-                                  SmallVectorImpl<Value>& outerIndices,
-                                  ConversionPatternRewriter& rewriter,
-                                  Location loc) {
+static FailureOr<Value> buildLoopSoftmaxNest(Value input,
+                                             Value accumulator,
+                                             RankedTensorType inputType,
+                                             int64_t axis,
+                                             SmallVectorImpl<Value>& outerIndices,
+                                             ConversionPatternRewriter& rewriter,
+                                             Location loc) {
   if (axis == inputType.getRank() - 1)
     return buildLoopSoftmaxSlice(input, accumulator, inputType, outerIndices, rewriter, loc);
 
@@ -57,38 +58,50 @@ static Value buildLoopSoftmaxNest(Value input,
   Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
   Value cUpper = getOrCreateIndexConstant(rewriter, anchorOp, inputType.getDimSize(axis));
 
-  auto loop = scf::ForOp::create(rewriter, loc, c0, cUpper, c1, ValueRange {accumulator});
-  rewriter.setInsertionPointToStart(loop.getBody());
-
-  Value loopIndex = loop.getInductionVar();
-  Value loopAccumulator = loop.getRegionIterArgs().front();
-  outerIndices.push_back(loopIndex);
-  Value updatedAccumulator =
-    buildLoopSoftmaxNest(input, loopAccumulator, inputType, axis + 1, outerIndices, rewriter, loc);
-  outerIndices.pop_back();
-
-  scf::YieldOp::create(rewriter, loc, updatedAccumulator);
-  rewriter.setInsertionPointAfter(loop);
-  return loop.getResult(0);
+  auto loop = buildNormalizedScfFor(
+    rewriter,
+    loc,
+    c0,
+    cUpper,
+    c1,
+    ValueRange {accumulator},
+    [&](OpBuilder& builder, Location nestedLoc, Value loopIndex, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+      outerIndices.push_back(loopIndex);
+      auto updatedAccumulator =
+        buildLoopSoftmaxNest(input, iterArgs.front(), inputType, axis + 1, outerIndices, rewriter, nestedLoc);
+      outerIndices.pop_back();
+      if (failed(updatedAccumulator))
+        return failure();
+      yielded.push_back(*updatedAccumulator);
+      return success();
+    });
+  if (failed(loop))
+    return failure();
+  return loop->results.front();
 }
 
-static Value createLoopSoftmaxCompute(Value input, ConversionPatternRewriter& rewriter, Location loc) {
+static FailureOr<Value> createLoopSoftmaxCompute(Value input, ConversionPatternRewriter& rewriter, Location loc) {
   auto inputType = cast<RankedTensorType>(input.getType());
   constexpr size_t numInputs = 1;
-  auto computeOp =
-    createSpatCompute<numInputs>(rewriter, loc, TypeRange {inputType}, {}, ValueRange {input}, [&](Value x) {
+  auto computeOp = createSpatCompute<numInputs>(
+    rewriter, loc, TypeRange {inputType}, {}, ValueRange {input}, [&](Value x) -> LogicalResult {
       if (inputType.getRank() == 1) {
         Value softmax = spatial::SpatSoftmaxOp::create(rewriter, loc, inputType, x).getResult();
         spatial::SpatYieldOp::create(rewriter, loc, softmax);
-        return;
+        return success();
       }
 
       Value outputInit = tensor::EmptyOp::create(rewriter, loc, inputType.getShape(), inputType.getElementType());
       SmallVector<Value> outerIndices;
-      Value result = buildLoopSoftmaxNest(x, outputInit, inputType, /*axis=*/0, outerIndices, rewriter, loc);
-      spatial::SpatYieldOp::create(rewriter, loc, result);
+      auto result = buildLoopSoftmaxNest(x, outputInit, inputType, /*axis=*/0, outerIndices, rewriter, loc);
+      if (failed(result))
+        return failure();
+      spatial::SpatYieldOp::create(rewriter, loc, *result);
+      return success();
     });
-  return computeOp.getResult(0);
+  if (failed(computeOp))
+    return failure();
+  return computeOp->getResult(0);
 }
 
 struct SoftmaxToSpatialCompute : OpConversionPattern<ONNXSoftmaxOp> {
@@ -108,7 +121,10 @@ struct SoftmaxToSpatialCompute : OpConversionPattern<ONNXSoftmaxOp> {
     Value input = adaptor.getInput();
     Value result;
     if (*axis == inputType.getRank() - 1) {
-      result = createLoopSoftmaxCompute(input, rewriter, softmaxOp.getLoc());
+      auto computed = createLoopSoftmaxCompute(input, rewriter, softmaxOp.getLoc());
+      if (failed(computed))
+        return failure();
+      result = *computed;
     }
     else {
       SmallVector<int64_t> permutation;
@@ -122,8 +138,10 @@ struct SoftmaxToSpatialCompute : OpConversionPattern<ONNXSoftmaxOp> {
       auto transposedType = RankedTensorType::get(
         permuteShape(inputType.getShape(), permutation), inputType.getElementType(), inputType.getEncoding());
       Value transposedInput = transposeMaybeInCompute(input, transposedType, permutation, rewriter, softmaxOp.getLoc());
-      Value transposedResult = createLoopSoftmaxCompute(transposedInput, rewriter, softmaxOp.getLoc());
-      result = transposeMaybeInCompute(transposedResult, inputType, inversePermutation, rewriter, softmaxOp.getLoc());
+      auto transposedResult = createLoopSoftmaxCompute(transposedInput, rewriter, softmaxOp.getLoc());
+      if (failed(transposedResult))
+        return failure();
+      result = transposeMaybeInCompute(*transposedResult, inputType, inversePermutation, rewriter, softmaxOp.getLoc());
     }
 
     rewriter.replaceOp(softmaxOp, result);