Move into function and refactor

Author: Mr-Anyone

mlir/include/mlir/Conversion/LinalgToRock/LinalgToRock.h

@@ -14,6 +14,7 @@
 #define MLIR_CONVERSION_LINALGTOROCK_H
 
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
 
@@ -24,6 +25,13 @@ namespace mlir {
 namespace rock {
 void populateLinalgToRockConversionPattern(RewritePatternSet &pattern,
                                            MLIRContext *context);
+
+/// A tensor.insert_slice is said to be a rock.expand_stride if it satisfies the following:
+/// - dest is a tensor.empty with a single use
+/// - all offsets are zero
+/// - all strides are one
+/// - all slice sizes are static and match the source tensor shape
+bool isRockExpandStride(tensor::InsertSliceOp op);
 }
 } // namespace mlir
 

mlir/lib/Conversion/LinalgToRock/LinalgToRock.cpp

@@ -156,39 +156,38 @@ struct ExpandStrideConverter final
 };
 } // namespace
 
-LogicalResult ExpandStrideConverter::matchAndRewrite(
-    tensor::InsertSliceOp op, OpAdaptor adaptor,
-    ConversionPatternRewriter &rewriter) const {
-  /// The linalg-to-rock passes emits the following expression
-  /// for expanding the strides. We are matching the following IR
-  /// %empty = tensor.empty() : ....
-  /// %inserted_slice = tensor.insert_slice %actual_data into %empty ...
-  auto tensorEmpty =
-      dyn_cast<tensor::EmptyOp>(op.getOperand(1).getDefiningOp());
-  if (!tensorEmpty || !tensorEmpty->hasOneUse()) {
-    return failure();
+bool mlir::rock::isRockExpandStride(tensor::InsertSliceOp op){
+  auto emptyOp = op.getDest().getDefiningOp<tensor::EmptyOp>();
+  if (!emptyOp){
+    return false;
   }
 
   // Require statically known slice sizes that exactly match the
   // source tensor shape.
   auto srcType = dyn_cast<RankedTensorType>(op.getSource().getType());
-  if (!srcType){
-    return failure();
-  }
+  if (!srcType)
+    return false;
 
-  // into rock.expand_strides, but only in the exact expand-strides shape:
-  // - dest is a tensor.empty with a single use
-  // - all offsets are zero
-  // - all strides are one
-  // - all slice sizes are static and match the source tensor shape
   bool isExpandStride = llvm::all_of(op.getStaticOffsets(), [](int64_t offset) { return offset == 0; }) &&
     llvm::all_of(op.getStaticStrides(), [](int64_t stride) { return stride == 1; }) &&
     llvm::none_of(op.getStaticSizes(),
         [](int64_t s) { return s == ShapedType::kDynamic; }) && op.getStaticSizes() == srcType.getShape();
+  return isExpandStride;
+}
 
-  if(!isExpandStride){
+LogicalResult ExpandStrideConverter::matchAndRewrite(
+    tensor::InsertSliceOp op, OpAdaptor adaptor,
+    ConversionPatternRewriter &rewriter) const {
+  /// The linalg-to-rock passes emits the following expression
+  /// for expanding the strides. We are matching the following IR
+  /// %empty = tensor.empty() : ....
+  /// %inserted_slice = tensor.insert_slice %actual_data into %empty ...
+  if(!rock::isRockExpandStride(op)){
     return failure();
   }
+  auto tensorEmpty =
+      dyn_cast<tensor::EmptyOp>(op.getOperand(1).getDefiningOp());
+  assert(tensorEmpty && "Should have been checked by isRockExpandStride");
 
   Location loc = op.getLoc();
   auto alloc = bufferization::AllocTensorOp::create(

mlir/lib/Conversion/LinalgToRock/LinalgToRockPass.cpp

@@ -38,31 +38,11 @@ static void populateLinalgToRockDialectConversion(ConversionTarget &target) {
                          rock::RockDialect, bufferization::BufferizationDialect,
                          math::MathDialect>();
 
-  // tensor.insert_slice with operand of one tensor.empty op can be expanded
-  // into rock.expand_strides
+  // a tensor.insert_slice oculd be a rock expand stride, and in that case
+  // we expand it into a rock.expand_stride
   target.addDynamicallyLegalOp<tensor::InsertSliceOp>(
       [](tensor::InsertSliceOp op) -> std::optional<bool> {
-        auto emptyOp = op.getDest().getDefiningOp<tensor::EmptyOp>();
-        if (!emptyOp){
-          return true;
-        }
-
-        // Require statically known slice sizes that exactly match the
-        // source tensor shape.
-        auto srcType = dyn_cast<RankedTensorType>(op.getSource().getType());
-        if (!srcType)
-          return true;
-
-        // into rock.expand_strides, but only in the exact expand-strides shape:
-        // - dest is a tensor.empty with a single use
-        // - all offsets are zero
-        // - all strides are one
-        // - all slice sizes are static and match the source tensor shape
-        bool isExpandStride = llvm::all_of(op.getStaticOffsets(), [](int64_t offset) { return offset == 0; }) &&
-          llvm::all_of(op.getStaticStrides(), [](int64_t stride) { return stride == 1; }) &&
-          llvm::none_of(op.getStaticSizes(),
-                  [](int64_t s) { return s == ShapedType::kDynamic; }) && op.getStaticSizes() == srcType.getShape();
-        return !isExpandStride;
+        return !rock::isRockExpandStride(op); 
       });
 
   // We only allow Linalg operations that are elementwise. Fusion is supported