[Mlir-commits] [mlir] [mlir][x86vector] Shuffle FMAs (PR #172823)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Dec 18 08:08:46 PST 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Arun Thangamani (arun-thmn)
<details>
<summary>Changes</summary>
This patch Shuffles FMAs with x86vector operations as operands such that FMAs are grouped with respect to odd/even packed index.
Continuation to PR: https://github.com/llvm/llvm-project/pull/170267 to manage register allocation efficiently.
---
Patch is 22.16 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/172823.diff
6 Files Affected:
- (modified) mlir/include/mlir/Dialect/X86Vector/TransformOps/X86VectorTransformOps.td (+11)
- (modified) mlir/include/mlir/Dialect/X86Vector/Transforms.h (+4)
- (modified) mlir/lib/Dialect/X86Vector/TransformOps/X86VectorTransformOps.cpp (+5)
- (modified) mlir/lib/Dialect/X86Vector/Transforms/CMakeLists.txt (+1)
- (added) mlir/lib/Dialect/X86Vector/Transforms/ShuffleVectorFMAOps.cpp (+182)
- (added) mlir/test/Dialect/X86Vector/shuffle-vector-fmas.mlir (+309)
``````````diff
diff --git a/mlir/include/mlir/Dialect/X86Vector/TransformOps/X86VectorTransformOps.td b/mlir/include/mlir/Dialect/X86Vector/TransformOps/X86VectorTransformOps.td
index 12ba5e9f11141..1183d1bb8b9f3 100644
--- a/mlir/include/mlir/Dialect/X86Vector/TransformOps/X86VectorTransformOps.td
+++ b/mlir/include/mlir/Dialect/X86Vector/TransformOps/X86VectorTransformOps.td
@@ -49,6 +49,17 @@ def ApplySinkVectorProducerOpsPatternsOp : Op<Transform_Dialect,
let assemblyFormat = "attr-dict";
}
+def ApplyShuffleVectorFMAOpsPatternsOp : Op<Transform_Dialect,
+ "apply_patterns.x86vector.shuffle_vector_fma_ops",
+ [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+ let description = [{
+ Collect patterns to shuffle FMAs with x86vector operations as operands
+ such that FMAs are grouped with respect to odd/even packed index.
+ }];
+
+ let assemblyFormat = "attr-dict";
+}
+
#endif // X86VECTOR_TRANSFORM_OPS
diff --git a/mlir/include/mlir/Dialect/X86Vector/Transforms.h b/mlir/include/mlir/Dialect/X86Vector/Transforms.h
index b9c9054f57890..72f3f685d7ec1 100644
--- a/mlir/include/mlir/Dialect/X86Vector/Transforms.h
+++ b/mlir/include/mlir/Dialect/X86Vector/Transforms.h
@@ -95,6 +95,10 @@ void populateVectorContractToPackedTypeDotProductPatterns(
// range by placing them at their earliest legal use site
void populateSinkVectorProducerOpsPatterns(RewritePatternSet &patterns);
+// Shuffles FMAs with x86vector operations as operands such that FMAs are
+// grouped with respect to odd/even packed index.
+void populateShuffleVectorFMAOpsPatterns(RewritePatternSet &patterns);
+
//===----------------------------------------------------------------------===//
/// Helpers extracted from:
/// - clang/lib/Headers/avxintrin.h
diff --git a/mlir/lib/Dialect/X86Vector/TransformOps/X86VectorTransformOps.cpp b/mlir/lib/Dialect/X86Vector/TransformOps/X86VectorTransformOps.cpp
index 25772f2aa57f4..a23f9b701199e 100644
--- a/mlir/lib/Dialect/X86Vector/TransformOps/X86VectorTransformOps.cpp
+++ b/mlir/lib/Dialect/X86Vector/TransformOps/X86VectorTransformOps.cpp
@@ -37,6 +37,11 @@ void mlir::transform::ApplySinkVectorProducerOpsPatternsOp::populatePatterns(
x86vector::populateSinkVectorProducerOpsPatterns(patterns);
}
+void mlir::transform::ApplyShuffleVectorFMAOpsPatternsOp::populatePatterns(
+ RewritePatternSet &patterns) {
+ x86vector::populateShuffleVectorFMAOpsPatterns(patterns);
+}
+
//===----------------------------------------------------------------------===//
// Transform op registration
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/X86Vector/Transforms/CMakeLists.txt b/mlir/lib/Dialect/X86Vector/Transforms/CMakeLists.txt
index cc4d3cac0f7ea..95471bc72f65d 100644
--- a/mlir/lib/Dialect/X86Vector/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/X86Vector/Transforms/CMakeLists.txt
@@ -4,6 +4,7 @@ add_mlir_dialect_library(MLIRX86VectorTransforms
VectorContractToFMA.cpp
VectorContractToPackedTypeDotProduct.cpp
SinkVectorProducerOps.cpp
+ ShuffleVectorFMAOps.cpp
LINK_LIBS PUBLIC
MLIRArithDialect
diff --git a/mlir/lib/Dialect/X86Vector/Transforms/ShuffleVectorFMAOps.cpp b/mlir/lib/Dialect/X86Vector/Transforms/ShuffleVectorFMAOps.cpp
new file mode 100644
index 0000000000000..945cf4e36b044
--- /dev/null
+++ b/mlir/lib/Dialect/X86Vector/Transforms/ShuffleVectorFMAOps.cpp
@@ -0,0 +1,182 @@
+//===- ShuffleVectorFMAOps.cpp --------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/X86Vector/Transforms.h"
+#include "mlir/Dialect/X86Vector/X86VectorDialect.h"
+
+#include "mlir/IR/PatternMatch.h"
+
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace mlir::vector;
+using namespace mlir::x86vector;
+
+namespace {
+
+// Validates whether the given operation is an x86vector operation and has only
+// one consumer.
+static bool validateX86OpsHasOneUser(Value op) {
+ if (auto cvt = op.getDefiningOp<x86vector::CvtPackedEvenIndexedToF32Op>())
+ return cvt.getResult().hasOneUse();
+
+ if (auto bcst = op.getDefiningOp<x86vector::BcstToPackedF32Op>())
+ return bcst.getResult().hasOneUse();
+
+ return false;
+}
+
+// Validates the vector.fma operation on the following conditions:
+// (i) one of the lhs or rhs defining operation should be
+// CvtPackedEvenIndexedToF32Op, (ii) the lhs or rhs defining operation should be
+// an x86vector operation and has only one consumer, (iii) all oerations in same
+// block, and (iv) ths FMA has only one user.
+static bool validateVectorFMAOp(vector::FMAOp fmaOp) {
+ Value lhs = fmaOp.getLhs();
+ Value rhs = fmaOp.getRhs();
+
+ if (!isa<x86vector::CvtPackedEvenIndexedToF32Op>(lhs.getDefiningOp()) &&
+ !isa<x86vector::CvtPackedEvenIndexedToF32Op>(rhs.getDefiningOp()))
+ return false;
+
+ if (!validateX86OpsHasOneUser(lhs) || !validateX86OpsHasOneUser(rhs))
+ return false;
+
+ if (lhs.getDefiningOp()->getBlock() != rhs.getDefiningOp()->getBlock())
+ return false;
+
+ if (lhs.getDefiningOp()->getBlock() != fmaOp->getBlock())
+ return false;
+
+ if (!fmaOp.getResult().hasOneUse())
+ return false;
+
+ Operation *consumer = *fmaOp.getResult().getUsers().begin();
+ if (consumer->getBlock() != fmaOp->getBlock())
+ return false;
+
+ return true;
+}
+
+// Moves vector.fma along with the lhs and rhs defining operation before it's
+// comsumer. If the consumer is vector.ShapeCastOp and has only one user then
+// move before the consumer of vector.ShapeCastOp.
+// TODO: Move before first consumer, if there are multiple.
+static void moveFMA(vector::FMAOp fmaOp) {
+ Operation *consumer = *fmaOp.getResult().getUsers().begin();
+
+ if (auto shapeCastOp = dyn_cast<vector::ShapeCastOp>(consumer)) {
+ if (shapeCastOp.getResult().hasOneUse()) {
+ Operation *nxtConsumer = *shapeCastOp.getResult().getUsers().begin();
+ if (nxtConsumer->getBlock() == fmaOp->getBlock()) {
+ consumer = *shapeCastOp.getResult().getUsers().begin();
+ fmaOp.getLhs().getDefiningOp()->moveBefore(consumer);
+ fmaOp.getRhs().getDefiningOp()->moveBefore(consumer);
+ fmaOp->moveBefore(consumer);
+ shapeCastOp->moveBefore(consumer);
+ return;
+ }
+ }
+ }
+
+ fmaOp.getLhs().getDefiningOp()->moveBefore(consumer);
+ fmaOp.getRhs().getDefiningOp()->moveBefore(consumer);
+ fmaOp->moveBefore(consumer);
+ return;
+}
+
+// Shuffle FMAs with x86vector operations as operands such that
+// FMAs are grouped with respect to odd/even packed index.
+//
+// For example:
+// ```
+// %1 = x86vector.avx.bcst_to_f32.packed
+// %2 = x86vector.avx.cvt.packed.odd.indexed_to_f32
+// %3 = vector.fma %1, %2, %arg1
+// %4 = x86vector.avx.bcst_to_f32.packed
+// %5 = x86vector.avx.cvt.packed.even.indexed_to_f32
+// %6 = vector.fma %4, %5, %3
+// %7 = x86vector.avx.bcst_to_f32.packed
+// %8 = x86vector.avx.cvt.packed.odd.indexed_to_f32
+// %9 = vector.fma %7, %8, %arg2
+// %10 = x86vector.avx.bcst_to_f32.packed
+// %11 = x86vector.avx.cvt.packed.even.indexed_to_f32
+// %12 = vector.fma %10, %11, %9
+// yield %6, %12
+// ```
+// to
+// ```
+// %1 = x86vector.avx.bcst_to_f32.packed
+// %2 = x86vector.avx.cvt.packed.odd.indexed_to_f32
+// %3 = vector.fma %1, %2, %arg1
+// %4 = x86vector.avx.bcst_to_f32.packed
+// %5 = x86vector.avx.cvt.packed.odd.indexed_to_f32
+// %6 = vector.fma %4, %5, %arg2
+// %7 = x86vector.avx.bcst_to_f32.packed
+// %8 = x86vector.avx.cvt.packed.even.indexed_to_f32
+// %9 = vector.fma %7, %8, %3
+// %10 = x86vector.avx.bcst_to_f32.packed
+// %11 = x86vector.avx.cvt.packed.even.indexed_to_f32
+// %12 = vector.fma %10, %11, %6
+// yield %9, %12
+// ```
+// TODO: Shuffling supported only if the FMA, lhs/rhs defining operations
+// have only one consumer. Have to extend this pass for multiple consumers.
+struct ShuffleVectorFMAOps : public OpRewritePattern<vector::FMAOp> {
+ using OpRewritePattern<vector::FMAOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(vector::FMAOp fmaOp,
+ PatternRewriter &rewriter) const override {
+
+ if (!validateVectorFMAOp(fmaOp))
+ return failure();
+
+ llvm::SmallVector<vector::FMAOp> fmaOps;
+ Operation *nextOp = fmaOp;
+ bool stopAtNextDependentFMA = true;
+
+ // Break the loop and return failure if the immediate next FMA op
+ // have CvtPackedEvenIndexedToF32Op in it's lhs/rhs defining ops.
+ while ((nextOp = nextOp->getNextNode())) {
+ auto fma = dyn_cast<vector::FMAOp>(nextOp);
+ if (!fma)
+ continue;
+
+ bool hasX86CvtOperand = isa<x86vector::CvtPackedEvenIndexedToF32Op>(
+ fma.getLhs().getDefiningOp()) ||
+ isa<x86vector::CvtPackedEvenIndexedToF32Op>(
+ fma.getRhs().getDefiningOp());
+
+ if (hasX86CvtOperand && stopAtNextDependentFMA)
+ break;
+
+ if (validateVectorFMAOp(fma))
+ fmaOps.push_back(fma);
+
+ stopAtNextDependentFMA = false;
+ }
+
+ if (fmaOps.empty())
+ return failure();
+
+ fmaOps.push_back(fmaOp);
+ for (auto fmaOp : fmaOps)
+ moveFMA(fmaOp);
+
+ return success();
+ }
+};
+
+} // namespace
+
+void x86vector::populateShuffleVectorFMAOpsPatterns(
+ RewritePatternSet &patterns) {
+ patterns.add<ShuffleVectorFMAOps>(patterns.getContext());
+}
diff --git a/mlir/test/Dialect/X86Vector/shuffle-vector-fmas.mlir b/mlir/test/Dialect/X86Vector/shuffle-vector-fmas.mlir
new file mode 100644
index 0000000000000..0e3c0b53b9cbd
--- /dev/null
+++ b/mlir/test/Dialect/X86Vector/shuffle-vector-fmas.mlir
@@ -0,0 +1,309 @@
+// RUN: mlir-opt %s -transform-interpreter -cse -split-input-file | FileCheck %s
+
+!vec = vector<8xf32>
+!memrefA = memref<1x1x1xbf16>
+!memrefB = memref<1x8x2xbf16>
+
+func.func @shuffle_fma_lhs_even_index(
+ %arg0: !memrefA, %arg1: !memrefA, %arg2: !memrefB, %arg3: !memrefA,
+ %arg4: !memrefA, %arg5: !memrefB, %arg6: !vec) -> !vec
+{
+ %0 = x86vector.avx.bcst_to_f32.packed %arg0 : !memrefA -> !vec
+ %1 = x86vector.avx.cvt.packed.odd.indexed_to_f32 %arg2 : !memrefB -> !vec
+ %2 = vector.fma %0, %1, %arg6 : !vec
+ %3 = x86vector.avx.bcst_to_f32.packed %arg1 : !memrefA -> !vec
+ %4 = x86vector.avx.cvt.packed.even.indexed_to_f32 %arg2 : !memrefB -> !vec
+ %5 = vector.fma %3, %4, %2 : !vec
+ %6 = x86vector.avx.bcst_to_f32.packed %arg3 : !memrefA -> !vec
+ %7 = x86vector.avx.cvt.packed.odd.indexed_to_f32 %arg5 : !memrefB -> !vec
+ %8 = vector.fma %6, %7, %arg6 : !vec
+ %9 = x86vector.avx.bcst_to_f32.packed %arg4 : !memrefA -> !vec
+ %10 = x86vector.avx.cvt.packed.even.indexed_to_f32 %arg5 : !memrefB -> !vec
+ %11 = vector.fma %9, %10, %8 : !vec
+ %12 = vector.fma %5, %11, %arg6 : !vec
+ return %12 : !vec
+}
+
+// CHECK-LABEL: @shuffle_fma_lhs_even_index
+// CHECK: x86vector.avx.bcst_to_f32.packed
+// CHECK: x86vector.avx.cvt.packed.odd.indexed_to_f32
+// CHECK: vector.fma
+// CHECK: x86vector.avx.bcst_to_f32.packed
+// CHECK: x86vector.avx.cvt.packed.odd.indexed_to_f32
+// CHECK: vector.fma
+// CHECK: x86vector.avx.bcst_to_f32.packed
+// CHECK: x86vector.avx.cvt.packed.even.indexed_to_f32
+// CHECK: vector.fma
+// CHECK: x86vector.avx.bcst_to_f32.packed
+// CHECK: x86vector.avx.cvt.packed.even.indexed_to_f32
+// CHECK: vector.fma
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %0 {
+ transform.apply_patterns.x86vector.shuffle_vector_fma_ops
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+!vec = vector<8xf32>
+!memrefA = memref<1x1x1xbf16>
+!memrefB = memref<1x8x2xbf16>
+
+func.func @shuffle_fma_rhs_even_index(
+ %arg0: !memrefA, %arg1: !memrefA, %arg2: !memrefB, %arg3: !memrefA,
+ %arg4: !memrefA, %arg5: !memrefB, %arg6: !vec) -> !vec
+{
+ %0 = x86vector.avx.bcst_to_f32.packed %arg0 : !memrefA -> !vec
+ %1 = x86vector.avx.cvt.packed.odd.indexed_to_f32 %arg2 : !memrefB -> !vec
+ %2 = vector.fma %0, %1, %arg6 : !vec
+ %3 = x86vector.avx.bcst_to_f32.packed %arg1 : !memrefA -> !vec
+ %4 = x86vector.avx.cvt.packed.even.indexed_to_f32 %arg2 : !memrefB -> !vec
+ %5 = vector.fma %4, %3, %2 : !vec
+ %6 = x86vector.avx.bcst_to_f32.packed %arg3 : !memrefA -> !vec
+ %7 = x86vector.avx.cvt.packed.odd.indexed_to_f32 %arg5 : !memrefB -> !vec
+ %8 = vector.fma %6, %7, %arg6 : !vec
+ %9 = x86vector.avx.bcst_to_f32.packed %arg4 : !memrefA -> !vec
+ %10 = x86vector.avx.cvt.packed.even.indexed_to_f32 %arg5 : !memrefB -> !vec
+ %11 = vector.fma %9, %10, %8 : !vec
+ %12 = vector.fma %5, %11, %arg6 : !vec
+ return %12 : !vec
+}
+
+// CHECK-LABEL: @shuffle_fma_rhs_even_index
+// CHECK: x86vector.avx.bcst_to_f32.packed
+// CHECK: x86vector.avx.cvt.packed.odd.indexed_to_f32
+// CHECK: vector.fma
+// CHECK: x86vector.avx.bcst_to_f32.packed
+// CHECK: x86vector.avx.cvt.packed.odd.indexed_to_f32
+// CHECK: vector.fma
+// CHECK: x86vector.avx.bcst_to_f32.packed
+// CHECK: x86vector.avx.cvt.packed.even.indexed_to_f32
+// CHECK: vector.fma
+// CHECK: x86vector.avx.cvt.packed.even.indexed_to_f32
+// CHECK: x86vector.avx.bcst_to_f32.packed
+// CHECK: vector.fma
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %0 {
+ transform.apply_patterns.x86vector.shuffle_vector_fma_ops
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+!vec = vector<8xf32>
+!memrefA = memref<1x1x1xbf16>
+!memrefB = memref<1x8x2xbf16>
+
+func.func @negative_fma_lhs_multiple_consumer(
+ %arg0: !memrefA, %arg1: !memrefA, %arg2: !memrefB,
+ %arg3: !memrefA, %arg4: !memrefB, %arg5: !vec) -> !vec
+{
+ %0 = x86vector.avx.bcst_to_f32.packed %arg0 : !memrefA -> !vec
+ %1 = x86vector.avx.cvt.packed.odd.indexed_to_f32 %arg2 : !memrefB -> !vec
+ %2 = vector.fma %0, %1, %arg5 : !vec
+ %3 = x86vector.avx.bcst_to_f32.packed %arg1 : !memrefA -> !vec
+ %4 = x86vector.avx.cvt.packed.even.indexed_to_f32 %arg2 : !memrefB -> !vec
+ %5 = vector.fma %3, %4, %2 : !vec
+ %7 = x86vector.avx.cvt.packed.odd.indexed_to_f32 %arg4 : !memrefB -> !vec
+ %8 = vector.fma %3, %7, %arg5 : !vec
+ %9 = x86vector.avx.bcst_to_f32.packed %arg3 : !memrefA -> !vec
+ %10 = x86vector.avx.cvt.packed.even.indexed_to_f32 %arg4 : !memrefB -> !vec
+ %11 = vector.fma %9, %10, %8 : !vec
+ %12 = vector.fma %5, %11, %arg5 : !vec
+ return %12 : !vec
+}
+
+// CHECK-LABEL: @negative_fma_lhs_multiple_consumer
+// CHECK: x86vector.avx.bcst_to_f32.packed
+// CHECK: x86vector.avx.cvt.packed.odd.indexed_to_f32
+// CHECK: vector.fma
+// CHECK: x86vector.avx.bcst_to_f32.packed
+// CHECK: x86vector.avx.cvt.packed.even.indexed_to_f32
+// CHECK: vector.fma
+// CHECK: x86vector.avx.cvt.packed.odd.indexed_to_f32
+// CHECK: vector.fma
+// CHECK: x86vector.avx.bcst_to_f32.packed
+// CHECK: x86vector.avx.cvt.packed.even.indexed_to_f32
+// CHECK: vector.fma
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %0 {
+ transform.apply_patterns.x86vector.shuffle_vector_fma_ops
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+!vec = vector<8xf32>
+!memrefA = memref<1x1x1xbf16>
+!memrefB = memref<1x8x2xbf16>
+
+func.func @negative_fma_rhs_multiple_consumer(
+ %arg0: !memrefA, %arg1: !memrefA, %arg2: !memrefB, %arg3: !memrefA,
+ %arg4: !memrefA, %arg5: !memrefB, %arg6: !vec) -> !vec
+{
+ %0 = x86vector.avx.bcst_to_f32.packed %arg0 : !memrefA -> !vec
+ %1 = x86vector.avx.cvt.packed.odd.indexed_to_f32 %arg2 : !memrefB -> !vec
+ %2 = vector.fma %0, %1, %arg6 : !vec
+ %3 = x86vector.avx.bcst_to_f32.packed %arg1 : !memrefA -> !vec
+ %4 = x86vector.avx.cvt.packed.even.indexed_to_f32 %arg2 : !memrefB -> !vec
+ %5 = vector.fma %3, %4, %2 : !vec
+ %6 = x86vector.avx.bcst_to_f32.packed %arg3 : !memrefA -> !vec
+ %7 = x86vector.avx.cvt.packed.odd.indexed_to_f32 %arg5 : !memrefB -> !vec
+ %8 = vector.fma %6, %7, %arg6 : !vec
+ %9 = x86vector.avx.bcst_to_f32.packed %arg4 : !memrefA -> !vec
+ %10 = vector.fma %9, %4, %8 : !vec
+ %11 = vector.fma %5, %10, %arg6 : !vec
+ return %11 : !vec
+}
+
+// CHECK-LABEL: @negative_fma_rhs_multiple_consumer
+// CHECK: x86vector.avx.bcst_to_f32.packed
+// CHECK: x86vector.avx.cvt.packed.odd.indexed_to_f32
+// CHECK: vector.fma
+// CHECK: x86vector.avx.bcst_to_f32.packed
+// CHECK: x86vector.avx.cvt.packed.even.indexed_to_f32
+// CHECK: vector.fma
+// CHECK: x86vector.avx.bcst_to_f32.packed
+// CHECK: x86vector.avx.cvt.packed.odd.indexed_to_f32
+// CHECK: vector.fma
+// CHECK: x86vector.avx.bcst_to_f32.packed
+// CHECK: vector.fma
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %0 {
+ transform.apply_patterns.x86vector.shuffle_vector_fma_ops
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+!vec = vector<8xf32>
+!memrefA = memref<1x1x1xbf16>
+!memrefB = memref<1x8x2xbf16>
+
+func.func @negative_fma_multiple_consumer(
+ %arg0: !memrefA, %arg1: !memrefA, %arg2: !memrefB, %arg3: !memrefA,
+ %arg4: !memrefA, %arg5: !memrefB, %arg6: !vec) -> !vec
+{
+ %0 = x86vector.avx.bcst_to_f32.packed %arg0 : !memrefA -> !vec
+ %1 = x86vector.avx.cvt.packed.odd.indexed_to_f32 %arg2 : !memrefB -> !vec
+ %2 = vector.fma %0, %1, %arg6 : !vec
+ %3 = x86vector.avx.bcst_to_f32.packed %arg1 : !memrefA -> !vec
+ %4 = x86vector.avx.cvt.packed.even.indexed_to_f32 %arg2 : !memrefB -> !vec
+ %5 = vector.fma %3, %4, %2 : !vec
+ %6 = x86vector.avx.bcst_to_f32.packed %arg3 : !memrefA -> !vec
+ %7 = x86vector.avx.cvt.packed.odd.indexed_to_f32 %arg5 : !memrefB -> !vec
+ %8 = vector.fma %6, %7, %5 : !vec
+ %9 = x86vector.avx.bcst_to_f32.packed %arg4 : !memrefA -> !vec
+ %10 = x86vector.avx.cvt.packed.even.indexed_to_f32 %arg5 : !memrefB -> !vec
+ %11 = vector.fma %9, %10, %8 : !vec
+ %12 = vector.fma %5, %11, %arg6 : !vec
+ return %12 : !vec
+}
+
+// CHECK-LABEL: @negative_fma_multiple_consumer
+// CHECK: x86vector.avx.bcst_to_f32.packed
+// CHECK: x86vector.avx.cvt.packed.odd.indexed_to_f32
+// CHECK: vector.fma
+// CHECK: x86vector.avx.bcst_to_f32.packed
+// CHECK: x86vector.avx.cvt.packed.even.indexed_to_f32
+// CHECK: vector.fma
+// CHECK: x86vector.avx.bcst_to_f32.packed
+// CHECK: x86vector.avx.cvt.packed.odd.indexed_to_f32
+// CHECK: vector.fma
+// CHECK: x86vector.avx.bcst_to_f32.packed
+// CHECK: x86vector.avx.cvt.packed.even.indexed_to_f32
+// CHECK: vector.fma
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %0 {
+ transform.apply_patterns.x86vector.shuffle_vector_fma_ops
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+!vec = vector<8xf32>
+!memrefA = memref<1x1x1xbf16>
+!memrefB = memref<1x8x2xbf16>
+
+func.func @negative_no_shuffle_outside_block(
+ %arg0: !memrefA, %arg1: !memrefA, %arg2: !memrefB, %arg3: !memrefA,
+ %arg4: !memrefA, %arg5: !memrefB, %arg6: !vec, %arg7: i1) -> !vec
+{
+ %0 = x86vector.avx.bcst_to_f32.packed %arg0 : !memrefA -> !vec
+ %1 = x86vector.avx.cvt.packed.odd.indexed_to_f32 %arg2 : !memref...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/172823
More information about the Mlir-commits
mailing list