[mlir][ArmSME] Add rewrites to swap extract of extend

mlir/lib/Dialect/ArmSME/Transforms/OuterProductFusion.cpp

@@ -261,6 +261,104 @@ class OuterProductFusion2Way
   }
 };
 
+// Rewrites: vector.extract(arith.extend) -> arith.extend(vector.extract).
+//
+// This transforms IR like:
+//   %0 = arith.extsi %src : vector<4x[8]xi8> to vector<4x[8]xi32>
+//   %1 = vector.extract %0[0] : vector<[8]xi32> from vector<4x[8]xi32>
+// Into:
+//   %0 = vector.extract %src[0] : vector<[8]xi8> from vector<4x[8]xi8>
+//   %1 = arith.extsi %0 : vector<[8]xi8> to vector<[8]xi32>
+//
+// This enables outer product fusion in the `-arm-sme-outer-product-fusion`
+// pass when the result is the input to an outer product.
+struct SwapVectorExtractOfArithExtend
+    : public OpRewritePattern<vector::ExtractOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::ExtractOp extractOp,
+                                PatternRewriter &rewriter) const override {
+    VectorType resultType = llvm::dyn_cast<VectorType>(extractOp.getType());
+    if (!resultType)
+      return rewriter.notifyMatchFailure(extractOp,
+                                         "extracted type is not a vector type");
+
+    auto numScalableDims = llvm::count(resultType.getScalableDims(), true);
+    if (numScalableDims != 1)
+      return rewriter.notifyMatchFailure(
+          extractOp, "extracted type is not a 1-D scalable vector type");
+
+    auto *extendOp = extractOp.getVector().getDefiningOp();
+    if (!isa_and_present<arith::ExtSIOp, arith::ExtUIOp, arith::ExtFOp>(
+            extendOp))
+      return rewriter.notifyMatchFailure(extractOp,
+                                         "extract not from extend op");
+
+    auto loc = extractOp.getLoc();
+    StringAttr extendOpName = extendOp->getName().getIdentifier();
+    Value extendSource = extendOp->getOperand(0);
+
+    // Create new extract from source of extend.
+    Value newExtract = rewriter.create<vector::ExtractOp>(
+        loc, extendSource, extractOp.getMixedPosition());
+
+    // Extend new extract to original result type.
+    Operation *newExtend =
+        rewriter.create(loc, extendOpName, Value(newExtract), resultType);
+
+    rewriter.replaceOp(extractOp, newExtend);
+
+    return success();
+  }
+};
+
+// Same as above, but for vector.scalable.extract.
+//
+// This transforms IR like:
+//   %0 = arith.extsi %src : vector<[8]xi8> to vector<[8]xi32>
+//   %1 = vector.scalable.extract %0[0] : vector<[4]xi32> from vector<[8]xi32>
+// Into:
+//   %0 = vector.scalable.extract %src[0] : vector<[4]xi8> from vector<[8]xi8>
+//   %1 = arith.extsi %0 : vector<[4]xi8> to vector<[4]xi32>
+//
+// This enables outer product fusion in the `-arm-sme-outer-product-fusion`
+// pass when the result is the input to an outer product.
+struct SwapVectorScalableExtractOfArithExtend
+    : public OpRewritePattern<vector::ScalableExtractOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::ScalableExtractOp extractOp,
+                                PatternRewriter &rewriter) const override {
+    auto *extendOp = extractOp.getSource().getDefiningOp();
+    if (!isa_and_present<arith::ExtSIOp, arith::ExtUIOp, arith::ExtFOp>(
+            extendOp))
+      return rewriter.notifyMatchFailure(extractOp,
+                                         "extract not from extend op");
+
+    auto loc = extractOp.getLoc();
+    VectorType resultType = extractOp.getResultVectorType();
+
+    Value extendSource = extendOp->getOperand(0);
+    StringAttr extendOpName = extendOp->getName().getIdentifier();
+    VectorType extendSourceVectorType =
+        cast<VectorType>(extendSource.getType());
+
+    // Create new extract from source of extend.
+    VectorType extractResultVectorType =
+        resultType.clone(extendSourceVectorType.getElementType());
+    Value newExtract = rewriter.create<vector::ScalableExtractOp>(
+        loc, extractResultVectorType, extendSource, extractOp.getPos());
+
+    // Extend new extract to original result type.
+    Operation *newExtend =
+        rewriter.create(loc, extendOpName, Value(newExtract), resultType);
+
+    rewriter.replaceOp(extractOp, newExtend);
+
+    return success();
+  }
+};
+
 struct OuterProductFusionPass
     : public arm_sme::impl::OuterProductFusionBase<OuterProductFusionPass> {
 
@@ -278,7 +376,11 @@ struct OuterProductFusionPass
 
 void mlir::arm_sme::populateOuterProductFusionPatterns(
     RewritePatternSet &patterns) {
-  patterns.add<OuterProductFusion2Way>(patterns.getContext());
+  MLIRContext *context = patterns.getContext();
+  // Note: High benefit to ensure extract(extend) are swapped first.
+  patterns.add<SwapVectorExtractOfArithExtend,
+               SwapVectorScalableExtractOfArithExtend>(context, 1024);
+  patterns.add<OuterProductFusion2Way>(context);
 }
 
 std::unique_ptr<Pass> mlir::arm_sme::createOuterProductFusionPass() {

mlir/test/Dialect/ArmSME/outer-product-fusion.mlir

@@ -213,6 +213,48 @@ func.func @outerproduct_sub_widening_2way_unsigned_i16i16i32(
   return %1 : vector<[4]x[4]xi32>
 }
 
+/// Tests for related patterns.
+
+// -----
+
+// CHECK-LABEL: @extract_from_arith_ext(
+// CHECK-SAME:                          %[[SRC:.*]]: vector<4x[8]xi8>
+// CHECK: %[[EXTRACT:.*]] = vector.extract %[[SRC]][0] : vector<[8]xi8> from vector<4x[8]xi8>
+// CHECK: %[[EXTEND:.*]] = arith.extsi %[[EXTRACT]] : vector<[8]xi8> to vector<[8]xi32>
+// CHECK: return %[[EXTEND]]
+func.func @extract_from_arith_ext(%src: vector<4x[8]xi8>) -> vector<[8]xi32> {
+  %0 = arith.extsi %src : vector<4x[8]xi8> to vector<4x[8]xi32>
+  %1 = vector.extract %0[0] : vector<[8]xi32> from vector<4x[8]xi32>
+  return %1 : vector<[8]xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @non_constant_extract_from_arith_ext(
+// CHECK-SAME:                                       %[[SRC:[a-z0-9]+]]: vector<4x[8]xi8>,
+// CHECK-SAME:                                       %[[DIM:[a-z0-9]+]]: index
+// CHECK: %[[EXTRACT:.*]] = vector.extract %[[SRC]][%[[DIM]]] : vector<[8]xi8> from vector<4x[8]xi8>
+// CHECK: %[[EXTEND:.*]] = arith.extui %[[EXTRACT]] : vector<[8]xi8> to vector<[8]xi32>
+// CHECK: return %[[EXTEND]]
+func.func @non_constant_extract_from_arith_ext(%src: vector<4x[8]xi8>, %dim: index) -> vector<[8]xi32> {
+  %0 = arith.extui %src : vector<4x[8]xi8> to vector<4x[8]xi32>
+  %1 = vector.extract %0[%dim] : vector<[8]xi32> from vector<4x[8]xi32>
+  return %1 : vector<[8]xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @scalable_extract_from_arith_ext(
+// CHECK-SAME:                                   %[[SRC:.*]]: vector<[8]xf16>
+// CHECK: %[[EXTRACT:.*]] = vector.scalable.extract %[[SRC]][0] : vector<[4]xf16> from vector<[8]xf16>
+// CHECK: %[[EXTEND:.*]] = arith.extf %[[EXTRACT]] : vector<[4]xf16> to vector<[4]xf32>
+// CHECK: return %[[EXTEND]]
+func.func @scalable_extract_from_arith_ext(%src: vector<[8]xf16>) -> vector<[4]xf32> {
+  %0 = arith.extf %src : vector<[8]xf16> to vector<[8]xf32>
+  %1 = vector.scalable.extract %0[0] : vector<[4]xf32> from vector<[8]xf32>
+  return %1 : vector<[4]xf32>
+}
+
 /// Negative tests
 
 // -----
@@ -362,3 +404,55 @@ func.func @outerproduct_widening_2way__bad_defining_op(
 
   return %1 : vector<[4]x[4]xf32>
 }
+
+/// Negative tests for related patterns.
+
+// -----
+
+/// Non-vector extracts should be ignored.
+
+// CHECK-LABEL: @extract_scalar_from_arith_ext
+// CHECK-NEXT: arith.extsi
+// CHECK-NEXT: vector.extract
+func.func @extract_scalar_from_arith_ext(%src: vector<4x[8]xi8>) -> i32 {
+  %0 = arith.extsi %src : vector<4x[8]xi8> to vector<4x[8]xi32>
+  %1 = vector.extract %0[0, 0] : i32 from vector<4x[8]xi32>
+  return %1 : i32
+}
+
+// -----
+
+/// Extracted type should be a 1-D scalable vector type.
+
+// CHECK-LABEL: @extract_fixed_1d_vec_from_arith_ext
+// CHECK-NEXT: arith.extsi
+// CHECK-NEXT: vector.extract
+func.func @extract_fixed_1d_vec_from_arith_ext(%src: vector<4x8xi8>) -> vector<8xi32> {
+  %0 = arith.extsi %src : vector<4x8xi8> to vector<4x8xi32>
+  %1 = vector.extract %0[0] : vector<8xi32> from vector<4x8xi32>
+  return %1 : vector<8xi32>
+}
+
+// -----
+
+/// Extract must come from an arith extend.
+
+// CHECK-LABEL: @extract_from_non_arith_ext
+// CHECK-NEXT: vector.extract
+// CHECK-NEXT: return
+func.func @extract_from_non_arith_ext(%src: vector<4x[8]xi32>) -> vector<[8]xi32> {
+  %0 = vector.extract %src[0] : vector<[8]xi32> from vector<4x[8]xi32>
+  return %0 : vector<[8]xi32>
+}
+
+// -----
+
+/// Scalable extract must come from an arith extend.
+
+// CHECK-LABEL: @scalable_extract_from_non_arith_ext
+// CHECK-NEXT: vector.scalable.extract
+// CHECK-NEXT: return
+func.func @scalable_extract_from_non_arith_ext(%src: vector<[8]xf32>) -> vector<[4]xf32> {
+  %0 = vector.scalable.extract %src[0] : vector<[4]xf32> from vector<[8]xf32>
+  return %0 : vector<[4]xf32>
+}