[mlir] Fix translation of splat constants to LLVM IR

[ftynse]

Summary

When converting splat constants for nested sequential LLVM IR types wrapped in
MLIR, the constant conversion was erroneously assuming it was always possible
to recursively construct a constant of a sequential type given only one value.
Instead, wait until all sequential types are unpacked recursively before
constructing a scalar constant and wrapping it into the surrounding sequential
type.

[merge_guards_bot]

Unit tests: pass. 61804 tests passed, 0 failed and 781 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[nicolasvasilache]

Thanks for fixing!

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

llvm::Constant *ModuleTranslation::getLLVMConstant(llvm::Type *llvmType,
   if (auto floatAttr = attr.dyn_cast<FloatAttr>())
     return llvm::ConstantFP::get(llvmType, floatAttr.getValue());
   if (auto funcAttr = attr.dyn_cast<FlatSymbolRefAttr>())
     return functionMapping.lookup(funcAttr.getValue());
   if (auto splatAttr = attr.dyn_cast<SplatElementsAttr>()) {
     auto *sequentialType = cast<llvm::SequentialType>(llvmType);
     auto elementType = sequentialType->getElementType();
     uint64_t numElements = sequentialType->getNumElements();
-    auto *child = getLLVMConstant(elementType, splatAttr.getSplatValue(), loc);
+    // 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.
+    llvm::Constant *child = getLLVMConstant(
+        elementType,
+        isa<llvm::SequentialType>(elementType) ? splatAttr
+                                               : splatAttr.getSplatValue(),
+        loc);
     if (llvmType->isVectorTy())
       return llvm::ConstantVector::getSplat(numElements, child);
     if (llvmType->isArrayTy()) {
       auto arrayType = llvm::ArrayType::get(elementType, numElements);
       SmallVector<llvm::Constant *, 8> constants(numElements, child);
       return llvm::ConstantArray::get(arrayType, constants);
     }
   }

mlir/test/Target/llvmir.mlir

 // CHECK-NEXT: %16 = lshr <4 x i64> %2, %2
   %13 = llvm.lshr %arg2, %arg2 : !llvm<"<4 x i64>">
 // CHECK-NEXT: %17 = ashr <4 x i64> %2, %2
   %14 = llvm.ashr %arg2, %arg2 : !llvm<"<4 x i64>">
 // CHECK-NEXT:    ret <4 x float> %4
   llvm.return %1 : !llvm<"<4 x float>">
 }
 
+// CHECK-LABEL: @vector_splat_1d
+llvm.func @vector_splat_1d() -> !llvm<"<4 x float>"> {
+  // CHECK: ret <4 x float> zeroinitializer
+  %0 = llvm.mlir.constant(dense<0.000000e+00> : vector<4xf32>) : !llvm<"<4 x float>">
+  llvm.return %0 : !llvm<"<4 x float>">
+}
+
+// CHECK-LABEL: @vector_splat_2d
+llvm.func @vector_splat_2d() -> !llvm<"[4 x <16 x float>]"> {
+  // CHECK: ret [4 x <16 x float>] zeroinitializer
+  %0 = llvm.mlir.constant(dense<0.000000e+00> : vector<4x16xf32>) : !llvm<"[4 x <16 x float>]">
+  llvm.return %0 : !llvm<"[4 x <16 x float>]">
+}
+
+// CHECK-LABEL: @vector_splat_3d
+llvm.func @vector_splat_3d() -> !llvm<"[4 x [16 x <4 x float>]]"> {
+  // CHECK: ret [4 x [16 x <4 x float>]] zeroinitializer
+  %0 = llvm.mlir.constant(dense<0.000000e+00> : vector<4x16x4xf32>) : !llvm<"[4 x [16 x <4 x float>]]">
+  llvm.return %0 : !llvm<"[4 x [16 x <4 x float>]]">
+}
+
+// CHECK-LABEL: @vector_splat_nonzero
+llvm.func @vector_splat_nonzero() -> !llvm<"<4 x float>"> {
+  // 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>) : !llvm<"<4 x float>">
+  llvm.return %0 : !llvm<"<4 x float>">
+}
+
 // CHECK-LABEL: @ops
 llvm.func @ops(%arg0: !llvm.float, %arg1: !llvm.float, %arg2: !llvm.i32, %arg3: !llvm.i32) -> !llvm<"{ float, i32 }"> {
 // CHECK-NEXT: fsub float %0, %1
   %0 = llvm.fsub %arg0, %arg1 : !llvm.float
 // CHECK-NEXT: %6 = sub i32 %2, %3
   %1 = llvm.sub %arg2, %arg3 : !llvm.i32
 // CHECK-NEXT: %7 = icmp slt i32 %2, %6
   %2 = llvm.icmp "slt" %arg2, %1 : !llvm.i32