[mlir][spirv] Implement SPIR-V lowering for `vector.deinterleave` #95313

angelz913 · 2024-06-12T21:20:30Z

Added a conversion for vector.deinterleave to the VectorToSPIRV pass.
Added LIT tests for the new conversion.

llvmbot · 2024-06-12T21:21:02Z

@llvm/pr-subscribers-mlir-spirv

@llvm/pr-subscribers-mlir

Author: Angel Zhang (angelz913)

Changes

Added a conversion for vector.deinterleave to the VectorToSPIRV pass.
Added LIT tests for the new conversion.

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

2 Files Affected:

(modified) mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp (+71-3)
(modified) mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir (+50)

diff --git a/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp b/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
index 92168cfa36147..b9a086cfc91a4 100644
--- a/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
+++ b/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
@@ -618,6 +618,74 @@ struct VectorInterleaveOpConvert final
   }
 };
 
+struct VectorDeinterleaveOpConvert final
+    : public OpConversionPattern<vector::DeinterleaveOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(vector::DeinterleaveOp deinterleaveOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+
+    // Check the result vector type.
+    VectorType oldResultType = deinterleaveOp.getResultVectorType();
+    Type newResultType = getTypeConverter()->convertType(oldResultType);
+    if (!newResultType)
+      return rewriter.notifyMatchFailure(deinterleaveOp,
+                                         "unsupported result vector type");
+
+    // Get location.
+    Location loc = deinterleaveOp->getLoc();
+
+    // Deinterleave the indices.
+    VectorType sourceType = deinterleaveOp.getSourceVectorType();
+    int n = sourceType.getNumElements();
+
+    // Output vectors of size 1 are converted to scalars by the type converter.
+    // We cannot use `spirv::VectorShuffleOp` directly in this case, and need to
+    // use `spirv::CompositeExtractOp`.
+    if (n == 2) {
+      spirv::CompositeExtractOp compositeExtractZero =
+          rewriter.create<spirv::CompositeExtractOp>(
+              loc, newResultType, adaptor.getSource(),
+              rewriter.getI32ArrayAttr({0}));
+
+      spirv::CompositeExtractOp compositeExtractOne =
+          rewriter.create<spirv::CompositeExtractOp>(
+              loc, newResultType, adaptor.getSource(),
+              rewriter.getI32ArrayAttr({1}));
+
+      rewriter.replaceOp(deinterleaveOp,
+                         {compositeExtractZero, compositeExtractOne});
+      return success();
+    }
+
+    // Indices for `res1`.
+    auto seqEven = llvm::seq<int64_t>(n / 2);
+    auto indicesEven =
+        llvm::map_to_vector(seqEven, [](int i) { return i * 2; });
+
+    // Indices for `res2`.
+    auto seqOdd = llvm::seq<int64_t>(n / 2);
+    auto indicesOdd =
+        llvm::map_to_vector(seqOdd, [](int i) { return i * 2 + 1; });
+
+    // Create two SPIR-V shuffles.
+    spirv::VectorShuffleOp shuffleEven =
+        rewriter.create<spirv::VectorShuffleOp>(
+            loc, newResultType, adaptor.getSource(), adaptor.getSource(),
+            rewriter.getI32ArrayAttr(indicesEven));
+
+    spirv::VectorShuffleOp shuffleOdd = rewriter.create<spirv::VectorShuffleOp>(
+        loc, newResultType, adaptor.getSource(), adaptor.getSource(),
+        rewriter.getI32ArrayAttr(indicesOdd));
+
+    // Replace deinterleaveOp with SPIR-V shuffles.
+    rewriter.replaceOp(deinterleaveOp, {shuffleEven, shuffleOdd});
+
+    return success();
+  }
+};
+
 struct VectorLoadOpConverter final
     : public OpConversionPattern<vector::LoadOp> {
   using OpConversionPattern::OpConversionPattern;
@@ -862,9 +930,9 @@ void mlir::populateVectorToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
       VectorReductionFloatMinMax<CL_FLOAT_MAX_MIN_OPS>,
       VectorReductionFloatMinMax<GL_FLOAT_MAX_MIN_OPS>, VectorShapeCast,
       VectorInsertStridedSliceOpConvert, VectorShuffleOpConvert,
-      VectorInterleaveOpConvert, VectorSplatPattern, VectorLoadOpConverter,
-      VectorStoreOpConverter>(typeConverter, patterns.getContext(),
-                              PatternBenefit(1));
+      VectorInterleaveOpConvert, VectorDeinterleaveOpConvert,
+      VectorSplatPattern, VectorLoadOpConverter, VectorStoreOpConverter>(
+      typeConverter, patterns.getContext(), PatternBenefit(1));
 
   // Make sure that the more specialized dot product pattern has higher benefit
   // than the generic one that extracts all elements.
diff --git a/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir b/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
index 2592d0fc04111..87823ab9afc0f 100644
--- a/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
+++ b/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
@@ -507,6 +507,56 @@ func.func @interleave_size1(%a: vector<1xf32>, %b: vector<1xf32>) -> vector<2xf3
 
 // -----
 
+// CHECK-LABEL: func @deinterleave_return0
+// CHECK-SAME: (%[[ARG0:.+]]: vector<4xf32>)
+//       CHECK: %[[SHUFFLE0:.*]] = spirv.VectorShuffle [0 : i32, 2 : i32] %[[ARG0]], %[[ARG0]] : vector<4xf32>, vector<4xf32> -> vector<2xf32>
+//       CHECK: %[[SHUFFLE1:.*]] = spirv.VectorShuffle [1 : i32, 3 : i32] %[[ARG0]], %[[ARG0]] : vector<4xf32>, vector<4xf32> -> vector<2xf32>
+//       CHECK: return %[[SHUFFLE0]]
+func.func @deinterleave_return0(%a: vector<4xf32>) -> vector<2xf32> {
+  %0, %1 = vector.deinterleave %a : vector<4xf32> -> vector<2xf32>
+  return %0 : vector<2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @deinterleave_return1
+// CHECK-SAME: (%[[ARG0:.+]]: vector<4xf32>)
+//       CHECK: %[[SHUFFLE0:.*]] = spirv.VectorShuffle [0 : i32, 2 : i32] %[[ARG0]], %[[ARG0]] : vector<4xf32>, vector<4xf32> -> vector<2xf32>
+//       CHECK: %[[SHUFFLE1:.*]] = spirv.VectorShuffle [1 : i32, 3 : i32] %[[ARG0]], %[[ARG0]] : vector<4xf32>, vector<4xf32> -> vector<2xf32>
+//       CHECK: return %[[SHUFFLE1]]
+func.func @deinterleave_return1(%a: vector<4xf32>) -> vector<2xf32> {
+  %0, %1 = vector.deinterleave %a : vector<4xf32> -> vector<2xf32>
+  return %1 : vector<2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @deinterleave_scalar_return0
+// CHECK-SAME: (%[[ARG0:.+]]: vector<2xf32>)
+//       CHECK: %[[EXTRACT0:.*]] = spirv.CompositeExtract %[[ARG0]][0 : i32] : vector<2xf32>
+//       CHECK: %[[EXTRACT1:.*]] = spirv.CompositeExtract %[[ARG0]][1 : i32] : vector<2xf32>
+//       CHECK: %[[RES:.*]] = builtin.unrealized_conversion_cast %[[EXTRACT0]] : f32 to vector<1xf32>
+//       CHECK: return %[[RES]]
+func.func @deinterleave_scalar_return0(%a: vector<2xf32>) -> vector<1xf32> {
+  %0, %1 = vector.deinterleave %a: vector<2xf32> -> vector<1xf32>
+  return %0 : vector<1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @deinterleave_scalar_return1
+// CHECK-SAME: (%[[ARG0:.+]]: vector<2xf32>)
+//       CHECK: %[[EXTRACT0:.*]] = spirv.CompositeExtract %[[ARG0]][0 : i32] : vector<2xf32>
+//       CHECK: %[[EXTRACT1:.*]] = spirv.CompositeExtract %[[ARG0]][1 : i32] : vector<2xf32>
+//       CHECK: %[[RES:.*]] = builtin.unrealized_conversion_cast %[[EXTRACT1]] : f32 to vector<1xf32>
+//       CHECK: return %[[RES]]
+func.func @deinterleave_scalar_return1(%a: vector<2xf32>) -> vector<1xf32> {
+  %0, %1 = vector.deinterleave %a: vector<2xf32> -> vector<1xf32>
+  return %1 : vector<1xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @reduction_add
 //  CHECK-SAME: (%[[V:.+]]: vector<4xi32>)
 //       CHECK:   %[[S0:.+]] = spirv.CompositeExtract %[[V]][0 : i32] : vector<4xi32>

kuhar

Looks good, just some nits

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir

github-actions · 2024-06-12T21:35:06Z

✅ With the latest revision this PR passed the C/C++ code formatter.

github-actions · 2024-06-12T21:35:20Z

⚠️ C/C++ code formatter, clang-format found issues in your code. ⚠️

You can test this locally with the following command:

git-clang-format --diff bf7c505847aa58af23f14ee986ee4bb7acf22e62 d1b940b86ab6259a4151f3a2ec961c649b7a953d -- mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

View the diff from clang-format here.

diff --git a/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp b/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
index 3cdd4ee524..58bfc7d280 100644
--- a/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
+++ b/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
@@ -648,10 +648,9 @@ struct VectorDeinterleaveOpConvert final
               loc, newResultType, adaptor.getSource(),
               rewriter.getI32ArrayAttr({0}));
 
-      auto elem1 =
-          rewriter.create<spirv::CompositeExtractOp>(
-              loc, newResultType, adaptor.getSource(),
-              rewriter.getI32ArrayAttr({1}));
+      auto elem1 = rewriter.create<spirv::CompositeExtractOp>(
+          loc, newResultType, adaptor.getSource(),
+          rewriter.getI32ArrayAttr({1}));
 
       rewriter.replaceOp(deinterleaveOp,
                          {compositeExtractZero, compositeExtractOne});

kuhar

LGTM, just 2 nits.

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

MacDue

LGTM, thanks for handling this!

1. Added a conversion for vector.deinterleave to the VectorToSPIRV pass. 2. Added LIT tests for the new conversion.

Co-authored-by: Jakub Kuderski <[email protected]>

This commit is dependent on #95313.

…vm#95313) 1. Added a conversion for `vector.deinterleave` to the `VectorToSPIRV` pass. 2. Added LIT tests for the new conversion. --------- Co-authored-by: Jakub Kuderski <[email protected]>

angelz913 requested review from antiagainst and kuhar as code owners June 12, 2024 21:20

llvmbot added mlir:spirv mlir labels Jun 12, 2024

kuhar reviewed Jun 12, 2024

View reviewed changes

angelz913 mentioned this pull request Jun 12, 2024

[mlir][vector] Support n-D vectors in i8 to i4 trunci emulation #94946

Merged

angelz913 force-pushed the vector-deinterleave-to-spirv branch from 997bb08 to ddf3125 Compare June 12, 2024 21:42

kuhar approved these changes Jun 13, 2024

View reviewed changes

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp Outdated Show resolved Hide resolved

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp Outdated Show resolved Hide resolved

MacDue approved these changes Jun 13, 2024

View reviewed changes

angelz913 mentioned this pull request Jun 13, 2024

[mlir][spirv] Add integration test for vector.deinterleave #95465

Closed

angelz913 and others added 12 commits June 13, 2024 20:21

[mlir][spirv] Implement SPIR-V lowering for vector.deinterleave

6bd9576

1. Added a conversion for vector.deinterleave to the VectorToSPIRV pass. 2. Added LIT tests for the new conversion.

Remove comment

9de11f1

Co-authored-by: Jakub Kuderski <[email protected]>

Fix naming style

f88cdf6

Co-authored-by: Jakub Kuderski <[email protected]>

Fix naming style

444aea2

Co-authored-by: Jakub Kuderski <[email protected]>

Fix naming style

c6142bd

Co-authored-by: Jakub Kuderski <[email protected]>

Remove empty line

b6c3310

Co-authored-by: Jakub Kuderski <[email protected]>

Remove comment

389eacc

Co-authored-by: Jakub Kuderski <[email protected]>

Refactor code

e72114a

Fix style

b7a00c2

Remove duplicate tests

9a3396c

Change comment

ff169d5

Co-authored-by: Jakub Kuderski <[email protected]>

Change comment

0335c7a

Co-authored-by: Jakub Kuderski <[email protected]>

angelz913 force-pushed the vector-deinterleave-to-spirv branch from e91f266 to 0335c7a Compare June 13, 2024 20:22

angelz913 mentioned this pull request Jun 13, 2024

[mlir][spirv] Add integration test for vector.deinterleave #95469

Merged

kuhar merged commit 597cde1 into llvm:main Jun 13, 2024
5 of 6 checks passed

kuhar pushed a commit that referenced this pull request Jun 14, 2024

[mlir][spirv] Add integration test for vector.deinterleave (#95469)

b1b7643

This commit is dependent on #95313.

angelz913 deleted the vector-deinterleave-to-spirv branch June 14, 2024 19:31

jpienaar mentioned this pull request Jun 17, 2024

[mlirc] Add missing extern C #95829

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][spirv] Implement SPIR-V lowering for `vector.deinterleave` #95313

[mlir][spirv] Implement SPIR-V lowering for `vector.deinterleave` #95313

Uh oh!

angelz913 commented Jun 12, 2024

Uh oh!

llvmbot commented Jun 12, 2024 •

edited

Loading

Uh oh!

kuhar left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented Jun 12, 2024 •

edited

Loading

Uh oh!

github-actions bot commented Jun 12, 2024

Uh oh!

kuhar left a comment

Uh oh!

Uh oh!

Uh oh!

MacDue left a comment

Uh oh!

Uh oh!

Uh oh!

[mlir][spirv] Implement SPIR-V lowering for vector.deinterleave #95313

[mlir][spirv] Implement SPIR-V lowering for vector.deinterleave #95313

Uh oh!

Conversation

angelz913 commented Jun 12, 2024

Uh oh!

llvmbot commented Jun 12, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kuhar left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented Jun 12, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Jun 12, 2024

Uh oh!

kuhar left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

MacDue left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

[mlir][spirv] Implement SPIR-V lowering for `vector.deinterleave` #95313

[mlir][spirv] Implement SPIR-V lowering for `vector.deinterleave` #95313

llvmbot commented Jun 12, 2024 •

edited

Loading

github-actions bot commented Jun 12, 2024 •

edited

Loading