blazingly faster

src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp

@@ -13,6 +13,7 @@
 #include <utility>
 
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
@@ -374,9 +375,6 @@ extractGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
   auto physicalType = mlir::dyn_cast<mlir::RankedTensorType>(physicalBatch.getType());
   if (!physicalType || physicalType.getRank() != fragmentType.getRank() + 1)
     return mlir::failure();
-  mlir::SmallVector<int64_t> selectedShape {1};
-  llvm::append_range(selectedShape, fragmentType.getShape());
-  auto selectedType = mlir::RankedTensorType::get(selectedShape, fragmentType.getElementType(), fragmentType.getEncoding());
   mlir::SmallVector<mlir::OpFoldResult> offsets {slot};
   mlir::SmallVector<mlir::OpFoldResult> sizes {rewriter.getIndexAttr(1)};
   mlir::SmallVector<mlir::OpFoldResult> strides {rewriter.getIndexAttr(1)};
@@ -385,11 +383,8 @@ extractGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
     sizes.push_back(rewriter.getIndexAttr(dim));
     strides.push_back(rewriter.getIndexAttr(1));
   }
-  mlir::Value selected = mlir::tensor::ExtractSliceOp::create(rewriter, loc, selectedType, physicalBatch, offsets, sizes, strides);
-  mlir::SmallVector<mlir::ReassociationIndices> reassociation {{0, 1}};
-  for (int64_t dim = 2; dim <= fragmentType.getRank(); ++dim)
-    reassociation.push_back({dim});
-  return mlir::tensor::CollapseShapeOp::create(rewriter, loc, fragmentType, selected, reassociation).getResult();
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, physicalBatch, fragmentType, {offsets, sizes, strides});
 }
 
 template <typename BodyFn>

src/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.cpp

@@ -3,6 +3,7 @@
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
@@ -38,6 +39,30 @@ Value createPaddedInputCompute(Value input,
   if (inputType == paddedInputType)
     return input;
 
+  auto producer = inputType.getRank() == 2 && paddedInputType.getRank() == 2
+    ? input.getDefiningOp<spatial::SpatGraphComputeBatch>()
+    : spatial::SpatGraphComputeBatch();
+  auto inputFragmentType = producer
+    ? spatial::getGraphBatchFragmentType(inputType, producer.getLaneCount())
+    : FailureOr<RankedTensorType>(failure());
+  auto paddedFragmentType = producer
+    ? spatial::getGraphBatchFragmentType(paddedInputType, producer.getLaneCount())
+    : FailureOr<RankedTensorType>(failure());
+  if (producer && succeeded(inputFragmentType) && succeeded(paddedFragmentType)) {
+    auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {paddedInputType}, producer.getLaneCount(), {}, input,
+      [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
+        auto fragment = extractGraphBatchPhysicalFragment(
+          rewriter, loc, args.inputs.front(), args.lane, *inputFragmentType);
+        if (failed(fragment))
+          return failure();
+        Value padded = createZeroPaddedTensor(*fragment, *paddedFragmentType, rewriter, loc);
+        publishGraphBatchPhysicalFragment(rewriter, loc, padded, args.outputs.front(), args.lane);
+        return success();
+      });
+    if (succeeded(batch))
+      return batch->getResult(0);
+  }
+
   auto computeOp = createSpatCompute<1>(rewriter, loc, TypeRange {paddedInputType}, {}, input, [&](Value computeInput) {
     Value paddedInput = createZeroPaddedTensor(computeInput, paddedInputType, rewriter, loc);
     spatial::SpatYieldOp::create(rewriter, loc, paddedInput);