[mlir][LLVM] Allow scalable vectors in ShuffleVectorOp

[jsetoain]

Summary

The current implementation of ShuffleVectorOp assumes all vectors are
scalable. LLVM IR allows shufflevector operations on scalable vectors,
and the current translation between LLVM Dialect and LLVM IR does the
rigth thing when the shuffle mask is all zeroes. This is required to
do a splat operation on a scalable vector, but it doesn't make sense
for scalable vectors outside of that operation, i.e.: with non-all zero
masks.

[jsetoain]

This is a WiP for now. This patch comes with an issues I'd like to discuss (and fix) before it gets approved. The discussion thread is: Llvm.shufflevector operations for scalable vectors

[jsetoain]

Rebase on main

[jsetoain]

Add constraint for shuffles on scalable vectors

[aartbik]

looks good, except the error reported in verify

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

1934

I don't quite understand the error text (starting with "on").
Can you please rephrase?

[jsetoain]

Change error message

[jsetoain]

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

1934

In the end, the output would be something like "llvm.shufflevector` op on scalable vectors...", but you are right, I shouldn't abuse error handling like that. Phrasing still feels a bit awkward to me, but I hope it's clearer this way.

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

 // Printing/parsing for LLVM::ShuffleVectorOp.
 //===----------------------------------------------------------------------===//
 // Expects vector to be of wrapped LLVM vector type and position to be of
 // wrapped LLVM i32 type.
 void LLVM::ShuffleVectorOp::build(OpBuilder &b, OperationState &result,
                                   Value v1, Value v2, ArrayAttr mask,
                                   ArrayRef<NamedAttribute> attrs) {
   auto containerType = v1.getType();
-  auto vType = LLVM::getFixedVectorType(
-      LLVM::getVectorElementType(containerType), mask.size());
+  auto vType = LLVM::getVectorType(
+      LLVM::getVectorElementType(containerType), mask.size(),
+      containerType.cast<VectorType>().isScalable());
   build(b, result, vType, v1, v2, mask);
   result.addAttributes(attrs);
 }
 
 void ShuffleVectorOp::print(OpAsmPrinter &p) {
   p << ' ' << getV1() << ", " << getV2() << " " << getMask();
   p.printOptionalAttrDict((*this)->getAttrs(), {"mask"});
   p << " : " << getV1().getType() << ", " << getV2().getType();
if (parser.getCurrentLocation(&loc) || parser.parseOperand(v1) ||
       parser.parseColonType(typeV1) || parser.parseComma() ||
       parser.parseType(typeV2) ||
       parser.resolveOperand(v1, typeV1, result.operands) ||
       parser.resolveOperand(v2, typeV2, result.operands))
     return failure();
   if (!LLVM::isCompatibleVectorType(typeV1))
     return parser.emitError(
         loc, "expected LLVM IR dialect vector type for operand #1");
-  auto vType = LLVM::getFixedVectorType(LLVM::getVectorElementType(typeV1),
-                                        maskAttr.size());
+  auto vType =
+      LLVM::getVectorType(LLVM::getVectorElementType(typeV1), maskAttr.size(),
+                          typeV1.cast<VectorType>().isScalable());
   result.addTypes(vType);
   return success();
 }
 
 LogicalResult ShuffleVectorOp::verify() {
   Type type1 = getV1().getType();
   Type type2 = getV2().getType();
   if (LLVM::getVectorElementType(type1) != LLVM::getVectorElementType(type2))
     return emitOpError("expected matching LLVM IR Dialect element types");
+  if (LLVM::isScalableVectorType(type1))
+    if (llvm::any_of(getMask(), [](Attribute attr) {
+          return attr.cast<IntegerAttr>().getInt() != 0;
+        }))
+      return emitOpError("expected a splat operation for scalable vectors");
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Implementations for LLVM::LLVMFuncOp.
 //===----------------------------------------------------------------------===//
 
 // Add the entry block to the function.

mlir/test/Dialect/LLVMIR/invalid.mlir

 
 // -----
 
 func @gep_out_of_bounds(%ptr: !llvm.ptr<struct<(i32, struct<(i32, f32)>)>>, %idx: i64) {
   // expected-error @below {{index 2 indexing a struct is out of bounds}}
   llvm.getelementptr %ptr[%idx, 1, 3] : (!llvm.ptr<struct<(i32, struct<(i32, f32)>)>>, i64) -> !llvm.ptr<i32>
   return
 }
+
+// -----
+
+func @non_splat_shuffle_on_scalable_vector(%arg0: vector<[4]xf32>) {
+  // expected-error@+1 {{expected a splat operation for scalable vectors}}
+  %0 = llvm.shufflevector %arg0, %arg0 [0 : i32, 0 : i32, 0 : i32, 1 : i32] : vector<[4]xf32>, vector<[4]xf32>
+  return
+}
 No newline at end of file

mlir/test/Dialect/LLVMIR/roundtrip.mlir

// CHECK:  = llvm.insertelement {{.*}} : vector<4xf32>
   %1 = llvm.insertelement %arg2, %arg0[%arg1 : i32] : vector<4xf32>
 // CHECK:  = llvm.shufflevector {{.*}} [0 : i32, 0 : i32, 0 : i32, 0 : i32, 7 : i32] : vector<4xf32>, vector<4xf32>
   %2 = llvm.shufflevector %arg0, %arg0 [0 : i32, 0 : i32, 0 : i32, 0 : i32, 7 : i32] : vector<4xf32>, vector<4xf32>
 // CHECK:  = llvm.mlir.constant(dense<1.000000e+00> : vector<4xf32>) : vector<4xf32>
   %3 = llvm.mlir.constant(dense<1.0> : vector<4xf32>) : vector<4xf32>
   return
 }
 
+// CHECK-LABEL: @scalable_vect
+func @scalable_vect(%arg0: vector<[4]xf32>, %arg1: i32, %arg2: f32) {
+// CHECK:  = llvm.extractelement {{.*}} : vector<[4]xf32>
+  %0 = llvm.extractelement %arg0[%arg1 : i32] : vector<[4]xf32>
+// CHECK:  = llvm.insertelement {{.*}} : vector<[4]xf32>
+  %1 = llvm.insertelement %arg2, %arg0[%arg1 : i32] : vector<[4]xf32>
+// CHECK:  = llvm.shufflevector {{.*}} [0 : i32, 0 : i32, 0 : i32, 0 : i32] : vector<[4]xf32>, vector<[4]xf32>
+  %2 = llvm.shufflevector %arg0, %arg0 [0 : i32, 0 : i32, 0 : i32, 0 : i32] : vector<[4]xf32>, vector<[4]xf32>
+// CHECK:  = llvm.mlir.constant(dense<1.000000e+00> : vector<[4]xf32>) : vector<[4]xf32>
+  %3 = llvm.mlir.constant(dense<1.0> : vector<[4]xf32>) : vector<[4]xf32>
+  return
+}
+
 // CHECK-LABEL: @alloca
 func @alloca(%size : i64) {
   // CHECK: llvm.alloca %{{.*}} x i32 : (i64) -> !llvm.ptr<i32>
   llvm.alloca %size x i32 {alignment = 0} : (i64) -> (!llvm.ptr<i32>)
   // CHECK: llvm.alloca %{{.*}} x i32 {alignment = 8 : i64} : (i64) -> !llvm.ptr<i32>
   llvm.alloca %size x i32 {alignment = 8} : (i64) -> (!llvm.ptr<i32>)
   llvm.return
 }

mlir/test/Target/LLVMIR/llvmir.mlir

 llvm.func @vect_i64idx(%arg0: vector<4xf32>, %arg1: i64, %arg2: f32) {
   // CHECK-NEXT: extractelement <4 x float> {{.*}}, i64
   // CHECK-NEXT: insertelement <4 x float> {{.*}}, float %2, i64
   %0 = llvm.extractelement %arg0[%arg1 : i64] : vector<4xf32>
   %1 = llvm.insertelement %arg2, %arg0[%arg1 : i64] : vector<4xf32>
   llvm.return
 }
 
+// CHECK-LABEL: @scalable_vect
+llvm.func @scalable_vect(%arg0: vector<[4]xf32>, %arg1: i32, %arg2: f32) {
+  // CHECK-NEXT: extractelement <vscale x 4 x float> {{.*}}, i32
+  // CHECK-NEXT: insertelement <vscale x 4 x float> {{.*}}, float %2, i32
+  // CHECK-NEXT: shufflevector <vscale x 4 x float> %0, <vscale x 4 x float> %0, <vscale x 4 x i32> zeroinitializer
+  %0 = llvm.extractelement %arg0[%arg1 : i32] : vector<[4]xf32>
+  %1 = llvm.insertelement %arg2, %arg0[%arg1 : i32] : vector<[4]xf32>
+  %2 = llvm.shufflevector %arg0, %arg0 [0 : i32, 0 : i32, 0 : i32, 0 : i32] : vector<[4]xf32>, vector<[4]xf32>
+  llvm.return
+}
+
+// CHECK-LABEL: @scalable_vect_i64idx
+llvm.func @scalable_vect_i64idx(%arg0: vector<[4]xf32>, %arg1: i64, %arg2: f32) {
+  // CHECK-NEXT: extractelement <vscale x 4 x float> {{.*}}, i64
+  // CHECK-NEXT: insertelement <vscale x 4 x float> {{.*}}, float %2, i64
+  %0 = llvm.extractelement %arg0[%arg1 : i64] : vector<[4]xf32>
+  %1 = llvm.insertelement %arg2, %arg0[%arg1 : i64] : vector<[4]xf32>
+  llvm.return
+}
+
 // CHECK-LABEL: @alloca
 llvm.func @alloca(%size : i64) {
   // Alignment automatically set by the LLVM IR builder when alignment attribute
   // is 0.
   //  CHECK: alloca {{.*}} align 4
   llvm.alloca %size x i32 {alignment = 0} : (i64) -> (!llvm.ptr<i32>)
   // CHECK-NEXT: alloca {{.*}} align 8
   llvm.alloca %size x i32 {alignment = 8} : (i64) -> (!llvm.ptr<i32>)