[mlir] support translation of multidimensional vectors to LLVM IR

[ftynse]

Summary

MLIR unlike LLVM IR supports multidimensional vector types. Such types are
lowered to nested LLVM IR arrays wrapping an LLVM IR vector for the innermost
dimension of the MLIR vector. MLIR supports constants of such types using
ElementsAttr for values. Introduce support for converting ElementsAttr into
LLVM IR Constant Aggregates recursively. This enables translation of
multidimensional vector constants from MLIR to LLVM IR.

[nicolasvasilache]

Great, thanks for pushing this forward Alex!

[merge_guards_bot]

Unit tests: pass. 61913 tests passed, 0 failed and 783 were skipped.

clang-tidy: unknown.

clang-format: pass.

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

[rriddle]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

35

This functions needs some documentation.

[ftynse]

Added in a922e23101b882e22fc14ffb78b5857954dc86f3

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Target/LLVMIR/ModuleTranslation.h"
 
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Module.h"
+#include "mlir/IR/StandardTypes.h"
 #include "mlir/Support/LLVM.h"
 
 #include "llvm/ADT/SetVector.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 
 using namespace mlir;
 using namespace mlir::LLVM;
 
+static llvm::Constant *
+buildSequentialConstant(ArrayRef<llvm::Constant *> &constants,
+                        ArrayRef<int64_t> shape, llvm::Type *type,
+                        Location loc) {
+  if (shape.empty()) {
+    llvm::Constant *result = constants.front();
+    constants = constants.drop_front();
+    return result;
+  }
+
+  if (!isa<llvm::SequentialType>(type)) {
+    emitError(loc) << "expected sequential LLVM types wrapping a scalar";
+    return nullptr;
+  }
+
+  llvm::Type *elementType = type->getSequentialElementType();
+  SmallVector<llvm::Constant *, 8> nested;
+  nested.reserve(shape.front());
+  for (int64_t i = 0; i < shape.front(); ++i) {
+    nested.push_back(buildSequentialConstant(constants, shape.drop_front(),
+                                             elementType, loc));
+    if (!nested.back())
+      return nullptr;
+  }
+
+  if (shape.size() == 1 && type->isVectorTy())
+    return llvm::ConstantVector::get(nested);
+  return llvm::ConstantArray::get(
+      llvm::ArrayType::get(elementType, shape.front()), nested);
+}
+
+static llvm::Type *getInnermostElementType(llvm::Type *type) {
+  while (isa<llvm::SequentialType>(type))
+    type = type->getSequentialElementType();
+  return type;
+}
+
 /// Create an LLVM IR constant of `llvmType` from the MLIR attribute `attr`.
 /// This currently supports integer, floating point, splat and dense element
 /// attributes and combinations thereof.  In case of error, report it to `loc`
 /// and return nullptr.
 llvm::Constant *ModuleTranslation::getLLVMConstant(llvm::Type *llvmType,
                                                    Attribute attr,
                                                    Location loc) {
   if (!attr)
     return llvm::UndefValue::get(llvmType);
+  if (llvmType->isStructTy()) {
+    emitError(loc, "struct types are not supported in constants");
+    return nullptr;
+  }
   if (auto intAttr = attr.dyn_cast<IntegerAttr>())
     return llvm::ConstantInt::get(llvmType, intAttr.getValue());
   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();
     // 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 (!child)
+      return nullptr;
     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);
     }
   }
+
   if (auto elementsAttr = attr.dyn_cast<ElementsAttr>()) {
-    auto *sequentialType = cast<llvm::SequentialType>(llvmType);
-    auto elementType = sequentialType->getElementType();
-    uint64_t numElements = sequentialType->getNumElements();
+    assert(elementsAttr.getType().hasStaticShape());
+    assert(elementsAttr.getNumElements() != 0 &&
+           "unexpected empty elements attribute");
+    assert(!elementsAttr.getType().getShape().empty() &&
+           "unexpected empty elements attribute shape");
+
     SmallVector<llvm::Constant *, 8> constants;
-    constants.reserve(numElements);
+    constants.reserve(elementsAttr.getNumElements());
+    llvm::Type *innermostType = getInnermostElementType(llvmType);
     for (auto n : elementsAttr.getValues<Attribute>()) {
-      constants.push_back(getLLVMConstant(elementType, n, loc));
+      constants.push_back(getLLVMConstant(innermostType, n, loc));
       if (!constants.back())
         return nullptr;
     }
-    if (llvmType->isVectorTy())
-      return llvm::ConstantVector::get(constants);
-    if (llvmType->isArrayTy()) {
-      auto arrayType = llvm::ArrayType::get(elementType, numElements);
-      return llvm::ConstantArray::get(arrayType, constants);
-    }
+    ArrayRef<llvm::Constant *> constantsRef = constants;
+    llvm::Constant *result = buildSequentialConstant(
+        constantsRef, elementsAttr.getType().getShape(), llvmType, loc);
+    assert(constantsRef.empty() && "did not consume all elemental constants");
+    return result;
   }
+
   if (auto stringAttr = attr.dyn_cast<StringAttr>()) {
     return llvm::ConstantDataArray::get(
         llvmModule->getContext(), ArrayRef<char>{stringAttr.getValue().data(),
                                                  stringAttr.getValue().size()});
   }
   emitError(loc, "unsupported constant value");
   return nullptr;
 }

mlir/test/Target/llvmir-invalid.mlir

-// RUN: mlir-translate -verify-diagnostics -mlir-to-llvmir %s
+// RUN: mlir-translate -verify-diagnostics -split-input-file -mlir-to-llvmir %s
 
 // expected-error @+1 {{unsupported module-level operation}}
 func @foo() {
   llvm.return
 }
+
+// -----
+
+llvm.func @no_nested_struct() -> !llvm<"[2 x [2 x [2 x {i32}]]]"> {
+  // expected-error @+1 {{struct types are not supported in constants}}
+  %0 = llvm.mlir.constant(dense<[[[1, 2], [3, 4]], [[42, 43], [44, 45]]]> : tensor<2x2x2xi32>) : !llvm<"[2 x [2 x [2 x {i32}]]]">
+  llvm.return %0 : !llvm<"[2 x [2 x [2 x {i32}]]]">
+}

mlir/test/Target/llvmir.mlir

 
 // Check that the auxiliary `null` operation is converted into a `null` value.
 // CHECK-LABEL: @null
 llvm.func @null() -> !llvm<"i32*"> {
   %0 = llvm.mlir.null : !llvm<"i32*">
   // CHECK: ret i32* null
   llvm.return %0 : !llvm<"i32*">
 }
+
+// Check that dense elements attributes are exported properly in constants.
+// CHECK-LABEL: @elements_constant_3d_vector
+llvm.func @elements_constant_3d_vector() -> !llvm<"[2 x [2 x <2 x i32>]]"> {
+  // CHECK: ret [2 x [2 x <2 x i32>]]
+  // CHECK-SAME: {{\[}}[2 x <2 x i32>] [<2 x i32> <i32 1, i32 2>, <2 x i32> <i32 3, i32 4>],
+  // CHECK-SAME:       [2 x <2 x i32>] [<2 x i32> <i32 42, i32 43>, <2 x i32> <i32 44, i32 45>]]
+  %0 = llvm.mlir.constant(dense<[[[1, 2], [3, 4]], [[42, 43], [44, 45]]]> : vector<2x2x2xi32>) : !llvm<"[2 x [2 x <2 x i32>]]">
+  llvm.return %0 : !llvm<"[2 x [2 x <2 x i32>]]">
+}
+
+// CHECK-LABEL: @elements_constant_3d_array
+llvm.func @elements_constant_3d_array() -> !llvm<"[2 x [2 x [2 x i32]]]"> {
+  // CHECK: ret [2 x [2 x [2 x i32]]]
+  // CHECK-SAME: {{\[}}[2 x [2 x i32]] {{\[}}[2 x i32] [i32 1, i32 2], [2 x i32] [i32 3, i32 4]],
+  // CHECK-SAME:       [2 x [2 x i32]] {{\[}}[2 x i32] [i32 42, i32 43], [2 x i32] [i32 44, i32 45]]]
+  %0 = llvm.mlir.constant(dense<[[[1, 2], [3, 4]], [[42, 43], [44, 45]]]> : tensor<2x2x2xi32>) : !llvm<"[2 x [2 x [2 x i32]]]">
+  llvm.return %0 : !llvm<"[2 x [2 x [2 x i32]]]">
+}