[mlir][Vector] Add vector.extract(vector.shuffle) folder #115105

dcaballe · 2024-11-06T01:24:49Z

This PR adds a folder for extracting an element from a vector shuffle. It turns something like:

   %shuffle = vector.shuffle %a, %b [0, 8, 7, 15]
     : vector<8xf32>, vector<8xf32>
   %extract = vector.extract %shuffle[3] : f32 from vector<4xf32>

into:

   %extract = vector.extract %b[7] : f32 from vector<8xf32>

This PR adds a folder for extracting an element from a vector shuffle. It turns something like: ``` %shuffle = vector.shuffle %a, %b [0, 8, 7, 15] : vector<8xf32>, vector<8xf32> %extract = vector.extract %shuffle[3] : f32 from vector<4xf32> ``` into: ``` %extract = vector.extract %b[7] : f32 from vector<4xf32> ```

llvmbot · 2024-11-06T01:25:25Z

@llvm/pr-subscribers-mlir

Author: Diego Caballero (dcaballe)

Changes

This PR adds a folder for extracting an element from a vector shuffle. It turns something like:

   %shuffle = vector.shuffle %a, %b [0, 8, 7, 15]
     : vector&lt;8xf32&gt;, vector&lt;8xf32&gt;
   %extract = vector.extract %shuffle[3] : f32 from vector&lt;4xf32&gt;

into:

   %extract = vector.extract %b[7] : f32 from vector&lt;4xf32&gt;

Full diff: https://github.com/llvm/llvm-project/pull/115105.diff

2 Files Affected:

(modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+42)
(modified) mlir/test/Dialect/Vector/canonicalize.mlir (+18)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index d8913251e56e9e..723044aa2b66e4 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -1705,6 +1705,46 @@ static Value foldExtractFromBroadcast(ExtractOp extractOp) {
   return extractOp.getResult();
 }
 
+/// Fold extractOp coming from ShuffleOp.
+///
+/// Example:
+///
+///   %shuffle = vector.shuffle %a, %b [0, 8, 7, 15]
+///     : vector<8xf32>, vector<8xf32>
+///   %extract = vector.extract %shuffle[3] : f32 from vector<4xf32>
+/// ->
+///   %extract = vector.extract %b[7] : f32 from vector<4xf32>
+///
+static Value foldExtractFromShuffle(ExtractOp extractOp) {
+  // TODO: Canonicalization for dynamic position not implemented yet.
+  if (extractOp.hasDynamicPosition())
+    return Value();
+
+  auto shuffleOp = extractOp.getVector().getDefiningOp<ShuffleOp>();
+  if (!shuffleOp)
+    return Value();
+
+  // TODO: 0-D or multi-dimensional vectors not supported yet.
+  if (shuffleOp.getResultVectorType().getRank() != 1)
+    return Value();
+
+  int64_t inputVecSize = shuffleOp.getV1().getType().getShape()[0];
+  auto shuffleMask = shuffleOp.getMask();
+  int64_t extractIdx = extractOp.getStaticPosition()[0];
+  int64_t shuffleIdx = shuffleMask[extractIdx];
+
+  // Find the shuffled vector to extract from based on the shuffle index.
+  if (shuffleIdx < inputVecSize) {
+    extractOp.setOperand(0, shuffleOp.getV1());
+    extractOp.setStaticPosition({shuffleIdx});
+  } else {
+    extractOp.setOperand(0, shuffleOp.getV2());
+    extractOp.setStaticPosition({shuffleIdx - inputVecSize});
+  }
+
+  return extractOp.getResult();
+}
+
 // Fold extractOp with source coming from ShapeCast op.
 static Value foldExtractFromShapeCast(ExtractOp extractOp) {
   // TODO: Canonicalization for dynamic position not implemented yet.
@@ -1953,6 +1993,8 @@ OpFoldResult ExtractOp::fold(FoldAdaptor) {
     return res;
   if (auto res = foldExtractFromBroadcast(*this))
     return res;
+  if (auto res = foldExtractFromShuffle(*this))
+    return res;
   if (auto res = foldExtractFromShapeCast(*this))
     return res;
   if (auto val = foldExtractFromExtractStrided(*this))
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index df87f86765a3a3..5ae769090dac66 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -740,6 +740,24 @@ func.func @fold_extract_broadcast(%a : vector<1xf32>) -> vector<8xf32> {
   %r = vector.extract %b[0] : vector<8xf32> from vector<1x8xf32>
   return %r : vector<8xf32>
 }
+// -----
+
+// CHECK-LABEL: @fold_extract_shuffle
+//  CHECK-SAME:   %[[A:.*]]: vector<8xf32>, %[[B:.*]]: vector<8xf32>
+//   CHECK-NOT:   vector.shuffle
+//       CHECK:   vector.extract %[[A]][0] : f32 from vector<8xf32>
+//       CHECK:   vector.extract %[[B]][0] : f32 from vector<8xf32>
+//       CHECK:   vector.extract %[[A]][7] : f32 from vector<8xf32>
+//       CHECK:   vector.extract %[[B]][7] : f32 from vector<8xf32>
+func.func @fold_extract_shuffle(%a : vector<8xf32>, %b : vector<8xf32>)
+                                -> (f32, f32, f32, f32) {
+  %shuffle = vector.shuffle %a, %b [0, 8, 7, 15] : vector<8xf32>, vector<8xf32>
+  %e0 = vector.extract %shuffle[0] : f32 from vector<4xf32>
+  %e1 = vector.extract %shuffle[1] : f32 from vector<4xf32>
+  %e2 = vector.extract %shuffle[2] : f32 from vector<4xf32>
+  %e3 = vector.extract %shuffle[3] : f32 from vector<4xf32>
+  return %e0, %e1, %e2, %e3 : f32, f32, f32, f32
+}
 
 // -----

llvmbot · 2024-11-06T01:25:25Z

@llvm/pr-subscribers-mlir-vector

Author: Diego Caballero (dcaballe)

Changes

This PR adds a folder for extracting an element from a vector shuffle. It turns something like:

   %shuffle = vector.shuffle %a, %b [0, 8, 7, 15]
     : vector&lt;8xf32&gt;, vector&lt;8xf32&gt;
   %extract = vector.extract %shuffle[3] : f32 from vector&lt;4xf32&gt;

into:

   %extract = vector.extract %b[7] : f32 from vector&lt;4xf32&gt;

Full diff: https://github.com/llvm/llvm-project/pull/115105.diff

2 Files Affected:

(modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+42)
(modified) mlir/test/Dialect/Vector/canonicalize.mlir (+18)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index d8913251e56e9e..723044aa2b66e4 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -1705,6 +1705,46 @@ static Value foldExtractFromBroadcast(ExtractOp extractOp) {
   return extractOp.getResult();
 }
 
+/// Fold extractOp coming from ShuffleOp.
+///
+/// Example:
+///
+///   %shuffle = vector.shuffle %a, %b [0, 8, 7, 15]
+///     : vector<8xf32>, vector<8xf32>
+///   %extract = vector.extract %shuffle[3] : f32 from vector<4xf32>
+/// ->
+///   %extract = vector.extract %b[7] : f32 from vector<4xf32>
+///
+static Value foldExtractFromShuffle(ExtractOp extractOp) {
+  // TODO: Canonicalization for dynamic position not implemented yet.
+  if (extractOp.hasDynamicPosition())
+    return Value();
+
+  auto shuffleOp = extractOp.getVector().getDefiningOp<ShuffleOp>();
+  if (!shuffleOp)
+    return Value();
+
+  // TODO: 0-D or multi-dimensional vectors not supported yet.
+  if (shuffleOp.getResultVectorType().getRank() != 1)
+    return Value();
+
+  int64_t inputVecSize = shuffleOp.getV1().getType().getShape()[0];
+  auto shuffleMask = shuffleOp.getMask();
+  int64_t extractIdx = extractOp.getStaticPosition()[0];
+  int64_t shuffleIdx = shuffleMask[extractIdx];
+
+  // Find the shuffled vector to extract from based on the shuffle index.
+  if (shuffleIdx < inputVecSize) {
+    extractOp.setOperand(0, shuffleOp.getV1());
+    extractOp.setStaticPosition({shuffleIdx});
+  } else {
+    extractOp.setOperand(0, shuffleOp.getV2());
+    extractOp.setStaticPosition({shuffleIdx - inputVecSize});
+  }
+
+  return extractOp.getResult();
+}
+
 // Fold extractOp with source coming from ShapeCast op.
 static Value foldExtractFromShapeCast(ExtractOp extractOp) {
   // TODO: Canonicalization for dynamic position not implemented yet.
@@ -1953,6 +1993,8 @@ OpFoldResult ExtractOp::fold(FoldAdaptor) {
     return res;
   if (auto res = foldExtractFromBroadcast(*this))
     return res;
+  if (auto res = foldExtractFromShuffle(*this))
+    return res;
   if (auto res = foldExtractFromShapeCast(*this))
     return res;
   if (auto val = foldExtractFromExtractStrided(*this))
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index df87f86765a3a3..5ae769090dac66 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -740,6 +740,24 @@ func.func @fold_extract_broadcast(%a : vector<1xf32>) -> vector<8xf32> {
   %r = vector.extract %b[0] : vector<8xf32> from vector<1x8xf32>
   return %r : vector<8xf32>
 }
+// -----
+
+// CHECK-LABEL: @fold_extract_shuffle
+//  CHECK-SAME:   %[[A:.*]]: vector<8xf32>, %[[B:.*]]: vector<8xf32>
+//   CHECK-NOT:   vector.shuffle
+//       CHECK:   vector.extract %[[A]][0] : f32 from vector<8xf32>
+//       CHECK:   vector.extract %[[B]][0] : f32 from vector<8xf32>
+//       CHECK:   vector.extract %[[A]][7] : f32 from vector<8xf32>
+//       CHECK:   vector.extract %[[B]][7] : f32 from vector<8xf32>
+func.func @fold_extract_shuffle(%a : vector<8xf32>, %b : vector<8xf32>)
+                                -> (f32, f32, f32, f32) {
+  %shuffle = vector.shuffle %a, %b [0, 8, 7, 15] : vector<8xf32>, vector<8xf32>
+  %e0 = vector.extract %shuffle[0] : f32 from vector<4xf32>
+  %e1 = vector.extract %shuffle[1] : f32 from vector<4xf32>
+  %e2 = vector.extract %shuffle[2] : f32 from vector<4xf32>
+  %e3 = vector.extract %shuffle[3] : f32 from vector<4xf32>
+  return %e0, %e1, %e2, %e3 : f32, f32, f32, f32
+}
 
 // -----

hanhanW

LGTM, I have a question because I'm seeing a TODO. (Please wait for at least a day to see if there are feedback from others.)

hanhanW · 2024-11-06T02:07:20Z

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

+  // TODO: Canonicalization for dynamic position not implemented yet.
+  if (extractOp.hasDynamicPosition())
+    return Value();


Out of curiosity, how do we support the dynamic case?

the only thing that comes to my mind is a series of compare and branches?

Yes, that's what I'm thinking, and I'm wondering if there are other approaches. This approach looks bad to me because it introduces branching in the IR. If we don't have a good plan, perhaps we should drop the TODO and make an explicit comment.

The TODO is inherited from similar canonicalizations. Not sure if branches but at least quite some instructions, including more vector extracts and select ops. Definitely not a good canonical form for dynamic cases.

I see, thanks! Yes, I also think that it is not a good canonical form. It's good that we're on the same page..

c-rhodes

one minor comment but otherwise LGTM cheers

c-rhodes · 2024-11-06T09:18:30Z

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

+///     : vector<8xf32>, vector<8xf32>
+///   %extract = vector.extract %shuffle[3] : f32 from vector<4xf32>
+/// ->
+///   %extract = vector.extract %b[7] : f32 from vector<4xf32>


Suggested change

/// %extract = vector.extract %b[7] : f32 from vector<4xf32>

/// %extract = vector.extract %b[7] : f32 from vector<8xf32>

c-rhodes · 2024-11-06T09:29:17Z

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

+  // TODO: Canonicalization for dynamic position not implemented yet.
+  if (extractOp.hasDynamicPosition())
+    return Value();


the only thing that comes to my mind is a series of compare and branches?

banach-space

LGTM, thanks!

dcaballe requested a review from c-rhodes November 6, 2024 01:24

dcaballe requested review from banach-space and nicolasvasilache as code owners November 6, 2024 01:24

llvmbot added mlir:vectorops mlir mlir:vector labels Nov 6, 2024

dcaballe requested a review from hanhanW November 6, 2024 01:25

hanhanW approved these changes Nov 6, 2024

View reviewed changes

c-rhodes approved these changes Nov 6, 2024

View reviewed changes

banach-space approved these changes Nov 6, 2024

View reviewed changes

dcaballe added 2 commits November 6, 2024 15:31

Feedback

4c814a2

Merge remote-tracking branch 'origin/main' into shuffle-extract-canon

261104f

dcaballe merged commit af5c471 into llvm:main Nov 7, 2024
6 of 7 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][Vector] Add vector.extract(vector.shuffle) folder #115105

[mlir][Vector] Add vector.extract(vector.shuffle) folder #115105

Uh oh!

dcaballe commented Nov 6, 2024 •

edited by CoTinker

Loading

Uh oh!

llvmbot commented Nov 6, 2024

Uh oh!

llvmbot commented Nov 6, 2024

Uh oh!

hanhanW left a comment

Uh oh!

hanhanW Nov 6, 2024

Uh oh!

c-rhodes Nov 6, 2024

Uh oh!

hanhanW Nov 6, 2024

Uh oh!

dcaballe Nov 6, 2024

Uh oh!

hanhanW Nov 7, 2024

Uh oh!

c-rhodes left a comment

Uh oh!

c-rhodes Nov 6, 2024

Uh oh!

c-rhodes Nov 6, 2024

Uh oh!

banach-space left a comment

Uh oh!

Uh oh!

Uh oh!

	/// %extract = vector.extract %b[7] : f32 from vector<4xf32>
	/// %extract = vector.extract %b[7] : f32 from vector<8xf32>

[mlir][Vector] Add vector.extract(vector.shuffle) folder #115105

[mlir][Vector] Add vector.extract(vector.shuffle) folder #115105

Uh oh!

Conversation

dcaballe commented Nov 6, 2024 • edited by CoTinker Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Nov 6, 2024

Uh oh!

llvmbot commented Nov 6, 2024

Uh oh!

hanhanW left a comment

Choose a reason for hiding this comment

Uh oh!

hanhanW Nov 6, 2024

Choose a reason for hiding this comment

Uh oh!

c-rhodes Nov 6, 2024

Choose a reason for hiding this comment

Uh oh!

hanhanW Nov 6, 2024

Choose a reason for hiding this comment

Uh oh!

dcaballe Nov 6, 2024

Choose a reason for hiding this comment

Uh oh!

hanhanW Nov 7, 2024

Choose a reason for hiding this comment

Uh oh!

c-rhodes left a comment

Choose a reason for hiding this comment

Uh oh!

c-rhodes Nov 6, 2024

Choose a reason for hiding this comment

Uh oh!

c-rhodes Nov 6, 2024

Choose a reason for hiding this comment

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

dcaballe commented Nov 6, 2024 •

edited by CoTinker

Loading