diff --git a/mlir/lib/Dialect/ArmSME/Transforms/OuterProductFusion.cpp b/mlir/lib/Dialect/ArmSME/Transforms/OuterProductFusion.cpp index fa55f9b7b31e..7dc2aacd7e5f 100644 --- a/mlir/lib/Dialect/ArmSME/Transforms/OuterProductFusion.cpp +++ b/mlir/lib/Dialect/ArmSME/Transforms/OuterProductFusion.cpp @@ -261,6 +261,104 @@ private: } }; +// 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 { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(vector::ExtractOp extractOp, + PatternRewriter &rewriter) const override { + VectorType resultType = llvm::dyn_cast(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( + 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( + 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 { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(vector::ScalableExtractOp extractOp, + PatternRewriter &rewriter) const override { + auto *extendOp = extractOp.getSource().getDefiningOp(); + if (!isa_and_present( + 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(extendSource.getType()); + + // Create new extract from source of extend. + VectorType extractResultVectorType = + resultType.clone(extendSourceVectorType.getElementType()); + Value newExtract = rewriter.create( + 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 { @@ -278,7 +376,11 @@ struct OuterProductFusionPass void mlir::arm_sme::populateOuterProductFusionPatterns( RewritePatternSet &patterns) { - patterns.add(patterns.getContext()); + MLIRContext *context = patterns.getContext(); + // Note: High benefit to ensure extract(extend) are swapped first. + patterns.add(context, 1024); + patterns.add(context); } std::unique_ptr mlir::arm_sme::createOuterProductFusionPass() { diff --git a/mlir/test/Dialect/ArmSME/outer-product-fusion.mlir b/mlir/test/Dialect/ArmSME/outer-product-fusion.mlir index 6a200848c2d8..f24943cac4f7 100644 --- a/mlir/test/Dialect/ArmSME/outer-product-fusion.mlir +++ b/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> +}