[MLIR][VectorToLLVM] Handle scalable dim in createVectorLengthValue() #93361

zhaoshiz · 2024-05-25T01:14:20Z

LLVM's Vector Predication Intrinsics require an explicit vector length parameter: https://llvm.org/docs/LangRef.html#vector-predication-intrinsics.

For a scalable vector type, this should be caculated as VectorScaleOp multiplied by base vector length, e.g.: for <[4]xf32> we should return: vscale * 4.

LLVM's Vector Predication Intrinsics require an explicit vector length parameter: https://llvm.org/docs/LangRef.html#vector-predication-intrinsics. For a scalable vector type, this should be caculated as VectorScaleOp multiplied by base vector length, e.g.: for <[4]xf32> we should return: vscale * 4.

llvmbot · 2024-05-25T01:14:50Z

@llvm/pr-subscribers-mlir

Author: Zhaoshi Zheng (zhaoshiz)

Changes

LLVM's Vector Predication Intrinsics require an explicit vector length parameter: https://llvm.org/docs/LangRef.html#vector-predication-intrinsics.

For a scalable vector type, this should be caculated as VectorScaleOp multiplied by base vector length, e.g.: for <[4]xf32> we should return: vscale * 4.

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

2 Files Affected:

(modified) mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp (+12-2)
(modified) mlir/test/Conversion/VectorToLLVM/vector-reduction-to-llvm.mlir (+38)

diff --git a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
index fe6bcc1c8b667..18bd9660525b4 100644
--- a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
+++ b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
@@ -523,7 +523,7 @@ static Value getOrCreateAccumulator(ConversionPatternRewriter &rewriter,
                                      llvmType);
 }
 
-/// Creates a constant value with the 1-D vector shape provided in `llvmType`.
+/// Creates a value with the 1-D vector shape provided in `llvmType`.
 /// This is used as effective vector length by some intrinsics supporting
 /// dynamic vector lengths at runtime.
 static Value createVectorLengthValue(ConversionPatternRewriter &rewriter,
@@ -532,9 +532,19 @@ static Value createVectorLengthValue(ConversionPatternRewriter &rewriter,
   auto vShape = vType.getShape();
   assert(vShape.size() == 1 && "Unexpected multi-dim vector type");
 
-  return rewriter.create<LLVM::ConstantOp>(
+  Value vLen = rewriter.create<LLVM::ConstantOp>(
       loc, rewriter.getI32Type(),
       rewriter.getIntegerAttr(rewriter.getI32Type(), vShape[0]));
+
+  if (!vType.getScalableDims()[0])
+    return vLen;
+
+  // Create VScale*vShape[0] and return it as vector length.
+  Value vScale = rewriter.create<vector::VectorScaleOp>(loc);
+  vScale = rewriter.create<arith::IndexCastOp>(
+      loc, rewriter.getI32Type(), vScale);
+  vLen = rewriter.create<arith::MulIOp>(loc, vLen, vScale);
+  return vLen;
 }
 
 /// Helper method to lower a `vector.reduction` op that performs an arithmetic
diff --git a/mlir/test/Conversion/VectorToLLVM/vector-reduction-to-llvm.mlir b/mlir/test/Conversion/VectorToLLVM/vector-reduction-to-llvm.mlir
index f98a05f8d17e2..209afa217437b 100644
--- a/mlir/test/Conversion/VectorToLLVM/vector-reduction-to-llvm.mlir
+++ b/mlir/test/Conversion/VectorToLLVM/vector-reduction-to-llvm.mlir
@@ -79,6 +79,25 @@ func.func @masked_reduce_add_f32(%arg0: vector<16xf32>, %mask : vector<16xi1>) -
 // CHECK:           "llvm.intr.vp.reduce.fadd"(%[[NEUTRAL]], %[[INPUT]], %[[MASK]], %[[VL]]) : (f32, vector<16xf32>, vector<16xi1>, i32) -> f32
 
 
+// -----
+
+func.func @masked_reduce_add_f32_scalable(%arg0: vector<[4]xf32>, %mask : vector<[4]xi1>) -> f32 {
+  %0 = vector.mask %mask { vector.reduction <add>, %arg0 : vector<[4]xf32> into f32 } : vector<[4]xi1> -> f32
+  return %0 : f32
+}
+
+// CHECK-LABEL:   func.func @masked_reduce_add_f32_scalable(
+// CHECK-SAME:                              %[[INPUT:.*]]: vector<[4]xf32>,
+// CHECK-SAME:                              %[[MASK:.*]]: vector<[4]xi1>) -> f32 {
+// CHECK:           %[[NEUTRAL:.*]] = llvm.mlir.constant(0.000000e+00 : f32) : f32
+// CHECK:           %[[VL_BASE:.*]] = llvm.mlir.constant(4 : i32) : i32
+// CHECK:           %[[VSCALE:.*]] = "llvm.intr.vscale"() : () -> i64
+// CHECK:           %[[CAST_IDX:.*]] = builtin.unrealized_conversion_cast %[[VSCALE]] : i64 to index
+// CHECK:           %[[CAST_I32:.*]] = arith.index_cast %[[CAST_IDX]] : index to i32
+// CHECK:           %[[VL_MUL:.*]] = arith.muli %[[VL_BASE]], %[[CAST_I32]] : i32
+// CHECK:           "llvm.intr.vp.reduce.fadd"(%[[NEUTRAL]], %[[INPUT]], %[[MASK]], %[[VL_MUL]]) : (f32, vector<[4]xf32>, vector<[4]xi1>, i32) -> f32
+
+
 // -----
 
 func.func @masked_reduce_mul_f32(%arg0: vector<16xf32>, %mask : vector<16xi1>) -> f32 {
@@ -167,6 +186,25 @@ func.func @masked_reduce_add_i8(%arg0: vector<32xi8>, %mask : vector<32xi1>) ->
 // CHECK:           "llvm.intr.vp.reduce.add"(%[[NEUTRAL]], %[[INPUT]], %[[MASK]], %[[VL]]) : (i8, vector<32xi8>, vector<32xi1>, i32) -> i8
 
 
+// -----
+
+func.func @masked_reduce_add_i8_scalable(%arg0: vector<[16]xi8>, %mask : vector<[16]xi1>) -> i8 {
+  %0 = vector.mask %mask { vector.reduction <add>, %arg0 : vector<[16]xi8> into i8 } : vector<[16]xi1> -> i8
+  return %0 : i8
+}
+
+// CHECK-LABEL:   func.func @masked_reduce_add_i8_scalable(
+// CHECK-SAME:                             %[[INPUT:.*]]: vector<[16]xi8>,
+// CHECK-SAME:                             %[[MASK:.*]]: vector<[16]xi1>) -> i8 {
+// CHECK:           %[[NEUTRAL:.*]] = llvm.mlir.constant(0 : i8) : i8
+// CHECK:           %[[VL_BASE:.*]] = llvm.mlir.constant(16 : i32) : i32
+// CHECK:           %[[VSCALE:.*]] = "llvm.intr.vscale"() : () -> i64
+// CHECK:           %[[CAST_IDX:.*]] = builtin.unrealized_conversion_cast %[[VSCALE]] : i64 to index
+// CHECK:           %[[CAST_I32:.*]] = arith.index_cast %[[CAST_IDX]] : index to i32
+// CHECK:           %[[VL_MUL:.*]] = arith.muli %[[VL_BASE]], %[[CAST_I32]] : i32
+// CHECK:           "llvm.intr.vp.reduce.add"(%[[NEUTRAL]], %[[INPUT]], %[[MASK]], %[[VL_MUL]]) : (i8, vector<[16]xi8>, vector<[16]xi1>, i32) -> i8
+
+
 // -----
 
 func.func @masked_reduce_mul_i8(%arg0: vector<32xi8>, %mask : vector<32xi1>) -> i8 {

github-actions · 2024-05-25T01:17:29Z

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

zhaoshiz · 2024-05-29T16:27:58Z

gentle ping...

banach-space

Makes sense, thanks. How did you decide what tests to "duplicate"? There seems to be more cases with vector.reduction.

banach-space · 2024-05-29T17:47:54Z

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

+  if (!vType.getScalableDims()[0])
+    return vLen;
+
+  // Create VScale*vShape[0] and return it as vector length.


[nit] We tend to write "vscale" rather than VScale. Also, why Shape rather than shape[0] or (even better, referring to a C++ variable): vShape[0]. In fact, you could rename vLen as baseVecLength to make the variable names more descriptive and use that in the comment.

Changed to vScale * baseVecLength, refering to actual variable names in code.

I feel duplicating all tests is a bit redundant.. the triton test case is using 'add', other vector.mask %mask {vector.reduction ...} tests don't offer additional coverage on code path in mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp, i.e., the part maps vector.reduction to llvm.intr.vp.reduce. is not changed. I've dup-ed some tests

zhaoshiz · 2024-06-05T16:09:36Z

gentle ping..

zhaoshiz · 2024-06-10T16:17:28Z

gentle ping again..

banach-space

Really sorry about the delay, Ive been a bit behind with reviews lately :(

One small comment, otherwise LG

banach-space · 2024-06-12T19:40:29Z

mlir/test/Conversion/VectorToLLVM/vector-reduction-to-llvm.mlir

@@ -79,6 +79,25 @@ func.func @masked_reduce_add_f32(%arg0: vector<16xf32>, %mask : vector<16xi1>) -
 // CHECK:           "llvm.intr.vp.reduce.fadd"(%[[NEUTRAL]], %[[INPUT]], %[[MASK]], %[[VL]]) : (f32, vector<16xf32>, vector<16xi1>, i32) -> f32


+// -----
+
+func.func @masked_reduce_add_f32_scalable(%arg0: vector<[4]xf32>, %mask : vector<[4]xi1>) -> f32 {


Please use identical shapes to what's used in @masked_reduce_add_f32. This way the only thing that changes is "scalability" rather than two things at a time. Same comment for other tests.

Thanks, I just updated the tests.

… counterparts of regular vectors.

banach-space

Lovely, thank you for working on this and apologies for the delay, LGTM!

…llvm#93361) LLVM's Vector Predication Intrinsics require an explicit vector length parameter: https://llvm.org/docs/LangRef.html#vector-predication-intrinsics. For a scalable vector type, this should be caculated as VectorScaleOp multiplied by base vector length, e.g.: for <[4]xf32> we should return: vscale * 4.

zhaoshiz requested review from banach-space, dcaballe and nicolasvasilache as code owners May 25, 2024 01:14

llvmbot added the mlir label May 25, 2024

Update per clang-format. NFC.

1addb3c

banach-space reviewed May 29, 2024

View reviewed changes

Improve readability and add more test cases. NFC.

6c588fc

banach-space reviewed Jun 12, 2024

View reviewed changes

Update test cases of scalable vectors to use the same shapes as their…

bdd5fa0

… counterparts of regular vectors.

banach-space approved these changes Jun 13, 2024

View reviewed changes

zhaoshiz merged commit abcbbe7 into llvm:main Jun 13, 2024
7 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[MLIR][VectorToLLVM] Handle scalable dim in createVectorLengthValue() #93361

[MLIR][VectorToLLVM] Handle scalable dim in createVectorLengthValue() #93361

Uh oh!

zhaoshiz commented May 25, 2024

Uh oh!

llvmbot commented May 25, 2024

Uh oh!

github-actions bot commented May 25, 2024 •

edited

Loading

Uh oh!

zhaoshiz commented May 29, 2024

Uh oh!

banach-space left a comment

Uh oh!

banach-space May 29, 2024

Uh oh!

zhaoshiz May 29, 2024

Uh oh!

zhaoshiz May 29, 2024

Uh oh!

zhaoshiz commented Jun 5, 2024

Uh oh!

zhaoshiz commented Jun 10, 2024

Uh oh!

banach-space left a comment

Uh oh!

banach-space Jun 12, 2024

Uh oh!

zhaoshiz Jun 12, 2024

Uh oh!

banach-space left a comment

Uh oh!

Uh oh!

Uh oh!

[MLIR][VectorToLLVM] Handle scalable dim in createVectorLengthValue() #93361

[MLIR][VectorToLLVM] Handle scalable dim in createVectorLengthValue() #93361

Uh oh!

Conversation

zhaoshiz commented May 25, 2024

Uh oh!

llvmbot commented May 25, 2024

Uh oh!

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

Uh oh!

Uh oh!

zhaoshiz commented May 29, 2024

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

banach-space May 29, 2024

Choose a reason for hiding this comment

Uh oh!

zhaoshiz May 29, 2024

Choose a reason for hiding this comment

Uh oh!

zhaoshiz May 29, 2024

Choose a reason for hiding this comment

Uh oh!

zhaoshiz commented Jun 5, 2024

Uh oh!

zhaoshiz commented Jun 10, 2024

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

banach-space Jun 12, 2024

Choose a reason for hiding this comment

Uh oh!

zhaoshiz Jun 12, 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!

github-actions bot commented May 25, 2024 •

edited

Loading