[mlir] Fix scalable type translation in splat element attr

[jsetoain]

Summary

LLVM Dialect Constant Op translations assume that if the attribute is a
vector, it's a fixed length one, generating an invalid translation for
constant scalable vector initializations.

[jsetoain]

Add missing test

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

if (auto floatAttr = attr.dyn_cast<FloatAttr>()) {
     return llvm::ConstantFP::get(llvmType, floatAttr.getValue());
   }
   if (auto funcAttr = attr.dyn_cast<FlatSymbolRefAttr>())
     return llvm::ConstantExpr::getBitCast(
         moduleTranslation.lookupFunction(funcAttr.getValue()), llvmType);
   if (auto splatAttr = attr.dyn_cast<SplatElementsAttr>()) {
     llvm::Type *elementType;
     uint64_t numElements;
+    bool isScalable = false;
     if (auto *arrayTy = dyn_cast<llvm::ArrayType>(llvmType)) {
       elementType = arrayTy->getElementType();
       numElements = arrayTy->getNumElements();
     } else if (auto *fVectorTy = dyn_cast<llvm::FixedVectorType>(llvmType)) {
       elementType = fVectorTy->getElementType();
       numElements = fVectorTy->getNumElements();
     } else if (auto *sVectorTy = dyn_cast<llvm::ScalableVectorType>(llvmType)) {
       elementType = sVectorTy->getElementType();
       numElements = sVectorTy->getMinNumElements();
+      isScalable = true;
     } else {
       llvm_unreachable("unrecognized constant vector type");
     }
     // Splat value is a scalar. Extract it only if the element type is not
     // another sequence type. The recursion terminates because each step removes
     // one outer sequential type.
     bool elementTypeSequential =
         isa<llvm::ArrayType, llvm::VectorType>(elementType);
     llvm::Constant *child = getLLVMConstant(
         elementType,
         elementTypeSequential ? splatAttr
                               : splatAttr.getSplatValue<Attribute>(),
         loc, moduleTranslation, false);
     if (!child)
       return nullptr;
     if (llvmType->isVectorTy())
       return llvm::ConstantVector::getSplat(
-          llvm::ElementCount::get(numElements, /*Scalable=*/false), child);
+          llvm::ElementCount::get(numElements, /*Scalable=*/isScalable), child);
     if (llvmType->isArrayTy()) {
       auto *arrayType = llvm::ArrayType::get(elementType, numElements);
       SmallVector<llvm::Constant *, 8> constants(numElements, child);
       return llvm::ConstantArray::get(arrayType, constants);
     }
   }
 
   // Try using raw elements data if possible.

mlir/test/Target/LLVMIR/llvmir-types.mlir

 //
 // Vectors.
 //
 
 // CHECK: declare <4 x i32> @return_v4_i32()
 llvm.func @return_v4_i32() -> vector<4xi32>
 // CHECK: declare <4 x float> @return_v4_float()
 llvm.func @return_v4_float() -> vector<4xf32>
+// CHECK: declare <vscale x 4 x float> @return_vs_4_float()
+llvm.func @return_vs_4_float() -> vector<[4]xf32>
 // CHECK: declare <vscale x 4 x i32> @return_vs_4_i32()
 llvm.func @return_vs_4_i32() -> !llvm.vec<?x4 x i32>
 // CHECK: declare <vscale x 8 x half> @return_vs_8_half()
 llvm.func @return_vs_8_half() -> !llvm.vec<?x8 x f16>
 // CHECK: declare <4 x i8*> @return_v_4_pi8()
 llvm.func @return_v_4_pi8() -> !llvm.vec<4xptr<i8>>
 
 //

mlir/test/Target/LLVMIR/llvmir.mlir

 
 // CHECK-LABEL: @vector_splat_1d
 llvm.func @vector_splat_1d() -> vector<4xf32> {
   // CHECK: ret <4 x float> zeroinitializer
   %0 = llvm.mlir.constant(dense<0.000000e+00> : vector<4xf32>) : vector<4xf32>
   llvm.return %0 : vector<4xf32>
 }
 
+// CHECK-LABEL: @vector_splat_1d_scalable
+llvm.func @vector_splat_1d_scalable() -> vector<[4]xf32> {
+  // CHECK: ret <vscale x 4 x float> zeroinitializer
+  %0 = llvm.mlir.constant(dense<0.000000e+00> : vector<[4]xf32>) : vector<[4]xf32>
+  llvm.return %0 : vector<[4]xf32>
+}
+
 // CHECK-LABEL: @vector_splat_2d
 llvm.func @vector_splat_2d() -> !llvm.array<4 x vector<16 x f32>> {
   // CHECK: ret [4 x <16 x float>] zeroinitializer
   %0 = llvm.mlir.constant(dense<0.000000e+00> : vector<4x16xf32>) : !llvm.array<4 x vector<16 x f32>>
   llvm.return %0 : !llvm.array<4 x vector<16 x f32>>
 }
 
 // CHECK-LABEL: @vector_splat_3d
 llvm.func @vector_splat_3d() -> !llvm.array<4 x array<16 x vector<4 x f32>>> {
   // CHECK: ret [4 x [16 x <4 x float>]] zeroinitializer
   %0 = llvm.mlir.constant(dense<0.000000e+00> : vector<4x16x4xf32>) : !llvm.array<4 x array<16 x vector<4 x f32>>>
   llvm.return %0 : !llvm.array<4 x array<16 x vector<4 x f32>>>
 }
 
 // CHECK-LABEL: @vector_splat_nonzero
 llvm.func @vector_splat_nonzero() -> vector<4xf32> {
   // CHECK: ret <4 x float> <float 1.000000e+00, float 1.000000e+00, float 1.000000e+00, float 1.000000e+00>
   %0 = llvm.mlir.constant(dense<1.000000e+00> : vector<4xf32>) : vector<4xf32>
   llvm.return %0 : vector<4xf32>
 }
 
+// CHECK-LABEL: @vector_splat_nonzero_scalable
+llvm.func @vector_splat_nonzero_scalable() -> vector<[4]xf32> {
+  // CHECK: ret <vscale x 4 x float> shufflevector (<vscale x 4 x float> insertelement (<vscale x 4 x float> poison, float 1.000000e+00, i32 0), <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer)
+  %0 = llvm.mlir.constant(dense<1.000000e+00> : vector<[4]xf32>) : vector<[4]xf32>
+  llvm.return %0 : vector<[4]xf32>
+}
+
 // CHECK-LABEL: @ops
 llvm.func @ops(%arg0: f32, %arg1: f32, %arg2: i32, %arg3: i32) -> !llvm.struct<(f32, i32)> {
 // CHECK-NEXT: fsub float %0, %1
   %0 = llvm.fsub %arg0, %arg1 : f32
 // CHECK-NEXT: %6 = sub i32 %2, %3
   %1 = llvm.sub %arg2, %arg3 : i32
 // CHECK-NEXT: %7 = icmp slt i32 %2, %6
   %2 = llvm.icmp "slt" %arg2, %1 : i32