Bose

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

@@ -2242,8 +2242,8 @@ static FailureOr<Value> rewriteInputKTiledConv(const ConvLoweringState& state,
                     rewriter, reduceLoc, paddedRowType, paddedPatchRow, aOffsets, aSizes, unitStrides);
                   SmallVector<OpFoldResult> bOffsets {kOffset, rewriter.getIndexAttr(0)};
                   SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(xbarDim), rewriter.getIndexAttr(xbarDim)};
-                  Value bTile = tensor::ExtractSliceOp::create(
-                    rewriter, reduceLoc, weightTileType, weightArg, bOffsets, bSizes, unitStrides);
+                  Value bTile = extractStaticSliceOrIdentity(
+                    rewriter, reduceLoc, weightArg, weightTileType, bOffsets, bSizes, unitStrides);
                   Value piece = spatial::SpatVMMOp::create(rewriter, reduceLoc, paddedRowType, bTile, aTile).getResult();
                   reduceYielded.push_back(
                     spatial::SpatVAddOp::create(rewriter, reduceLoc, paddedRowType, acc, piece).getResult());
@@ -2912,8 +2912,13 @@ static FailureOr<Value> createConvOutputFromRowStripHwc(Value inputHwc,
                 rewriter, reduceLoc, paddedRowType, paddedRow, aOffsets, aSizes, getUnitStrides(rewriter, 2));
               SmallVector<OpFoldResult> bOffsets {kOffset, rewriter.getIndexAttr(0)};
               SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(xbarDim), rewriter.getIndexAttr(xbarDim)};
-              Value bTile = tensor::ExtractSliceOp::create(
-                rewriter, reduceLoc, paddedWeightTileType, args.weights.front(), bOffsets, bSizes, getUnitStrides(rewriter, 2));
+              Value bTile = extractStaticSliceOrIdentity(rewriter,
+                                                         reduceLoc,
+                                                         args.weights.front(),
+                                                         paddedWeightTileType,
+                                                         bOffsets,
+                                                         bSizes,
+                                                         getUnitStrides(rewriter, 2));
               Value piece = spatial::SpatVMMOp::create(rewriter, reduceLoc, paddedRowType, bTile, aTile).getResult();
               reduceYielded.push_back(
                 spatial::SpatVAddOp::create(rewriter, reduceLoc, paddedRowType, reduceIterArgs.front(), piece).getResult());

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

@@ -285,9 +285,8 @@ static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
       SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(crossbarSize.getValue()),
                                         rewriter.getIndexAttr(crossbarSize.getValue())};
       SmallVector<OpFoldResult> unitStrides = getUnitStrides(rewriter, 2);
-      Value bTile =
-        tensor::ExtractSliceOp::create(rewriter, loc, bTileType, args.weights.front(), bOffsets, bSizes, unitStrides)
-          .getResult();
+      Value bTile = extractStaticSliceOrIdentity(
+        rewriter, loc, args.weights.front(), bTileType, bOffsets, bSizes, unitStrides);
       Value piece = spatial::SpatVMMOp::create(rewriter, loc, pieceType, bTile, aTile).getResult();
 
       SmallVector<OpFoldResult> pieceOffsets {args.lane, rewriter.getIndexAttr(0)};