[mlir][vector] Follow-up improvements for multi-dimensional vector.from_elements support #154664

yangtetris · 2025-08-21T03:39:58Z

This PR is a follow-up to #151175 that supported lowering multi-dimensional vector.from_elements op to LLVM by introducing a unrolling pattern.

Changes

Add `vector.shape_cast` based flattening pattern for `vector.from_elements`

This change introduces a new linearization pattern that uses vector.shape_cast to flatten multi-dimensional vector.from_elements operations. This provides an alternative approach to the unrolling-based method introduced in #151175.

Example:

// Before
%v = vector.from_elements %e0, %e1, %e2, %e3 : vector<2x2xf32>

// After
%flat = vector.from_elements %e0, %e1, %e2, %e3 : vector<4xf32>
%result = vector.shape_cast %flat : vector<4xf32> to vector<2x2xf32>

llvmbot · 2025-08-21T03:40:31Z

@llvm/pr-subscribers-mlir-spirv
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-vector

Author: Yang Bai (yangtetris)

Changes

This PR is a follow-up to #151175 that supported lowering multi-dimensional vector.from_elements op to LLVM by introducing a unrolling pattern.

Changes

1. Add `vector.shape_cast` based flattening pattern for `vector.from_elements`

This change introduces a new linearization pattern that uses vector.shape_cast to flatten multi-dimensional vector.from_elements operations. This provides an alternative approach to the unrolling-based method introduced in #151175.

Example:

// Before
%v = vector.from_elements %e0, %e1, %e2, %e3 : vector&lt;2x2xf32&gt;

// After
%flat = vector.from_elements %e0, %e1, %e2, %e3 : vector&lt;4xf32&gt;
%result = vector.shape_cast %flat : vector&lt;4xf32&gt; to vector&lt;2x2xf32&gt;

2. Integrate `UnrollFromElements` pattern into SPIR-V pipeline

The UnrollFromElements pattern from #151175 is now integrated into spirv::unrollVectorsInFuncBodies to ensure proper handling of multi-dimensional vectors in SPIR-V lowering. Additional tests are included to verify the pattern works correctly in the SPIR-V context.

Example:

// Before
%v = vector.from_elements %e0, %e1, %e2, %e3 : vector&lt;2x2xf32&gt;

// After test-spirv-vector-unrolling
%vec_1d_0 = vector.from_elements %e0, %e1 : vector&lt;2xf32&gt;
%vec_1d_1 = vector.from_elements %e2, %e3 : vector&lt;2xf32&gt;

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

5 Files Affected:

(modified) mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp (+3)
(modified) mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp (+38-1)
(modified) mlir/test/Conversion/ConvertToSPIRV/vector-unroll.mlir (+12)
(modified) mlir/test/Conversion/ConvertToSPIRV/vector.mlir (+11)
(modified) mlir/test/Dialect/Vector/linearize.mlir (+14)

diff --git a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
index 49f4ce8de7c76..7815c3b377316 100644
--- a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
+++ b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
@@ -1496,6 +1496,9 @@ LogicalResult mlir::spirv::unrollVectorsInFuncBodies(Operation *op) {
     auto options = vector::UnrollVectorOptions().setNativeShapeFn(
         [](auto op) { return mlir::spirv::getNativeVectorShape(op); });
     populateVectorUnrollPatterns(patterns, options);
+    // The unroll pattern for vector.from_elements op doesn't belong to the
+    // above set.
+    vector::populateVectorFromElementsLoweringPatterns(patterns);
     if (failed(applyPatternsGreedily(op, std::move(patterns))))
       return failure();
   }
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
index 491b448e9e1e9..2cb6d47f37128 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
@@ -762,6 +762,42 @@ struct LinearizeVectorStore final
   }
 };
 
+/// This pattern linearizes `vector.from_elements` operations by converting
+/// the result type to a 1-D vector while preserving all element values.
+/// The transformation creates a linearized `vector.from_elements` followed by
+/// a `vector.shape_cast` to restore the original multidimensional shape.
+///
+/// Example:
+///
+///     %0 = vector.from_elements %a, %b, %c, %d : vector<2x2xf32>
+///
+///   is converted to:
+///
+///     %0 = vector.from_elements %a, %b, %c, %d : vector<4xf32>
+///     %1 = vector.shape_cast %0 : vector<4xf32> to vector<2x2xf32>
+///
+struct LinearizeVectorFromElements final
+    : public OpConversionPattern<vector::FromElementsOp> {
+  using OpConversionPattern::OpConversionPattern;
+  LinearizeVectorFromElements(const TypeConverter &typeConverter,
+                              MLIRContext *context, PatternBenefit benefit = 1)
+      : OpConversionPattern(typeConverter, context, benefit) {}
+  LogicalResult
+  matchAndRewrite(vector::FromElementsOp fromElementsOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    VectorType dstTy =
+        getTypeConverter()->convertType<VectorType>(fromElementsOp.getType());
+    assert(dstTy && "vector type destination expected.");
+
+    auto elements = fromElementsOp.getElements();
+    assert(elements.size() == static_cast<size_t>(dstTy.getNumElements()) &&
+           "expected same number of elements");
+    rewriter.replaceOpWithNewOp<vector::FromElementsOp>(fromElementsOp, dstTy,
+                                                        elements);
+    return success();
+  }
+};
+
 } // namespace
 
 /// This method defines the set of operations that are linearizable, and hence
@@ -854,7 +890,8 @@ void mlir::vector::populateVectorLinearizeBasePatterns(
   patterns
       .add<LinearizeConstantLike, LinearizeVectorizable, LinearizeVectorBitCast,
            LinearizeVectorSplat, LinearizeVectorCreateMask, LinearizeVectorLoad,
-           LinearizeVectorStore>(typeConverter, patterns.getContext());
+           LinearizeVectorStore, LinearizeVectorFromElements>(
+          typeConverter, patterns.getContext());
 }
 
 void mlir::vector::populateVectorLinearizeShuffleLikeOpsPatterns(
diff --git a/mlir/test/Conversion/ConvertToSPIRV/vector-unroll.mlir b/mlir/test/Conversion/ConvertToSPIRV/vector-unroll.mlir
index d68ba44ee8840..48c6dd5dc8c31 100644
--- a/mlir/test/Conversion/ConvertToSPIRV/vector-unroll.mlir
+++ b/mlir/test/Conversion/ConvertToSPIRV/vector-unroll.mlir
@@ -100,3 +100,15 @@ func.func @transpose(%arg0 : vector<2x3xi32>) -> (vector<3x2xi32>) {
   %0 = vector.transpose %arg0, [1, 0] : vector<2x3xi32> to vector<3x2xi32>
   return %0 : vector<3x2xi32>
 }
+
+// -----
+
+// CHECK-LABEL: @from_elements_2d
+// CHECK-SAME: (%[[ARG0:.*]]: f32, %[[ARG1:.*]]: f32, %[[ARG2:.*]]: f32, %[[ARG3:.*]]: f32)
+func.func @from_elements_2d(%arg0: f32, %arg1: f32, %arg2: f32, %arg3: f32) -> vector<2x2xf32> {
+  // CHECK: %[[VEC0:.*]] = vector.from_elements %[[ARG0]], %[[ARG1]] : vector<2xf32>
+  // CHECK: %[[VEC1:.*]] = vector.from_elements %[[ARG2]], %[[ARG3]] : vector<2xf32>
+  // CHECK: return %[[VEC0]], %[[VEC1]] : vector<2xf32>, vector<2xf32>
+  %0 = vector.from_elements %arg0, %arg1, %arg2, %arg3 : vector<2x2xf32>
+  return %0 : vector<2x2xf32>
+}
diff --git a/mlir/test/Conversion/ConvertToSPIRV/vector.mlir b/mlir/test/Conversion/ConvertToSPIRV/vector.mlir
index eb9feaad15c5b..ff66fe576fc2a 100644
--- a/mlir/test/Conversion/ConvertToSPIRV/vector.mlir
+++ b/mlir/test/Conversion/ConvertToSPIRV/vector.mlir
@@ -356,3 +356,14 @@ func.func @shape_cast_size1_vector(%arg0 : vector<f32>) -> vector<1xf32> {
   %1 = vector.shape_cast %arg0 : vector<f32> to vector<1xf32>
   return %1 : vector<1xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @from_elements
+//  CHECK-SAME: (%[[ARG0:.*]]: f32, %[[ARG1:.*]]: f32, %[[ARG2:.*]]: f32)
+//       CHECK:   %[[RETVAL:.+]] = spirv.CompositeConstruct %[[ARG0]], %[[ARG1]], %[[ARG2]] : (f32, f32, f32) -> vector<3xf32>
+//       CHECK:   spirv.ReturnValue %[[RETVAL]]
+func.func @from_elements(%arg0: f32, %arg1: f32, %arg2: f32) -> vector<3xf32> {
+  %0 = vector.from_elements %arg0, %arg1, %arg2 : vector<3xf32>
+  return %0 : vector<3xf32>
+}
diff --git a/mlir/test/Dialect/Vector/linearize.mlir b/mlir/test/Dialect/Vector/linearize.mlir
index 2e630bf93622e..5e8bfd0698b33 100644
--- a/mlir/test/Dialect/Vector/linearize.mlir
+++ b/mlir/test/Dialect/Vector/linearize.mlir
@@ -524,3 +524,17 @@ func.func @linearize_vector_store_scalable(%arg0: memref<2x8xf32>, %arg1: vector
   vector.store %arg1, %arg0[%c0, %c0] : memref<2x8xf32>, vector<1x[4]xf32>
   return
 }
+
+// -----
+
+// Test pattern LinearizeVectorFromElements.
+
+// CHECK-LABEL: test_vector_from_elements
+// CHECK-SAME: %[[ARG_0:.*]]: f32, %[[ARG_1:.*]]: f32, %[[ARG_2:.*]]: f32, %[[ARG_3:.*]]: f32
+func.func @test_vector_from_elements(%arg0: f32, %arg1: f32, %arg2: f32, %arg3: f32) -> vector<2x2xf32> {
+  // CHECK: %[[FROM_ELEMENTS:.*]] = vector.from_elements %[[ARG_0]], %[[ARG_1]], %[[ARG_2]], %[[ARG_3]] : vector<4xf32>
+  // CHECK: %[[CAST:.*]] = vector.shape_cast %[[FROM_ELEMENTS]] : vector<4xf32> to vector<2x2xf32>
+  // CHECK: return %[[CAST]] : vector<2x2xf32>
+  %1 = vector.from_elements %arg0, %arg1, %arg2, %arg3 : vector<2x2xf32>
+  return %1 : vector<2x2xf32>
+}

newling

LGTM. FYI we've been discussing converting these patterns to greedy patterns #146030

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

Co-authored-by: James Newling <james.newling@gmail.com>

yangtetris · 2025-08-22T09:53:55Z

LGTM. FYI we've been discussing converting these patterns to greedy patterns #146030

Thanks. Actually, I was confused when I first saw that a type converter was used there.

dcaballe

Awesome, thanks!

banach-space

LGTM, thanks!

[mlir] support ND->1D flattening for vector.from_elements op

9471daa

yangtetris requested review from antiagainst, kuhar, banach-space, dcaballe, hanhanW and nicolasvasilache as code owners August 21, 2025 03:40

llvmbot added mlir:spirv mlir:vectorops mlir mlir:vector labels Aug 21, 2025

yangtetris marked this pull request as draft August 21, 2025 04:39

yangtetris force-pushed the mlir/vecor-from-elements-flattening branch from 3f76a9c to 9471daa Compare August 21, 2025 06:40

newling approved these changes Aug 22, 2025

View reviewed changes

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp Outdated Show resolved Hide resolved

newling reviewed Aug 22, 2025

View reviewed changes

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp Outdated Show resolved Hide resolved

yangtetris and others added 2 commits August 22, 2025 15:01

Update mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

5aa6bfa

Co-authored-by: James Newling <james.newling@gmail.com>

refine

4fe3766

dcaballe approved these changes Aug 22, 2025

View reviewed changes

kuhar requested a review from Groverkss August 23, 2025 06:03

yangtetris marked this pull request as ready for review August 26, 2025 05:02

banach-space approved these changes Aug 27, 2025

View reviewed changes

dcaballe merged commit 5fdd3a1 into llvm:main Aug 28, 2025
11 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][vector] Follow-up improvements for multi-dimensional vector.from_elements support #154664

[mlir][vector] Follow-up improvements for multi-dimensional vector.from_elements support #154664

yangtetris commented Aug 21, 2025 •

edited

Loading

Uh oh!

llvmbot commented Aug 21, 2025 •

edited

Loading

Changes

1. Add `vector.shape_cast` based flattening pattern for `vector.from_elements`

2. Integrate `UnrollFromElements` pattern into SPIR-V pipeline

Uh oh!

newling left a comment

Uh oh!

Uh oh!

Uh oh!

yangtetris commented Aug 22, 2025

Uh oh!

dcaballe left a comment

Uh oh!

banach-space left a comment

Uh oh!

Uh oh!

Uh oh!

[mlir][vector] Follow-up improvements for multi-dimensional vector.from_elements support #154664

[mlir][vector] Follow-up improvements for multi-dimensional vector.from_elements support #154664

Conversation

yangtetris commented Aug 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes

Add vector.shape_cast based flattening pattern for vector.from_elements

Uh oh!

llvmbot commented Aug 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes

1. Add vector.shape_cast based flattening pattern for vector.from_elements

2. Integrate UnrollFromElements pattern into SPIR-V pipeline

Uh oh!

newling left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

yangtetris commented Aug 22, 2025

Uh oh!

dcaballe left a comment

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!

yangtetris commented Aug 21, 2025 •

edited

Loading

Add `vector.shape_cast` based flattening pattern for `vector.from_elements`

llvmbot commented Aug 21, 2025 •

edited

Loading

1. Add `vector.shape_cast` based flattening pattern for `vector.from_elements`

2. Integrate `UnrollFromElements` pattern into SPIR-V pipeline