Add support for complex constants to MLIR core.

[akuegel]

I probably need to add more tests, and maybe need to improve some things. Still, I wanted to already send this out to get feedback on the general approach here.
The RFC with more discussion is here:
https://llvm.discourse.group/t/rfc-add-a-builtin-complexattr/3316/28

[rriddle]

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

1944

nit: Drop trivial braces here.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1024

1038

mlir/lib/IR/BuiltinAttributes.cpp

589

You shouldn't need vector here, can you not use ArrayRef directly, which supports creation from an initializer list? You might need to explicitly type that values as Attribute.

598

Same here.

[mehdi_amini]

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

1946

I think you also need to check more here, in particular that the two are the same type to form a valid complex?

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1045

Are the changes around standard constant exercised in a test?
Otherwise can you add this, including for the verifier errors added below.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

113

If you don't intend to support anything else than a complex, more invariants checking is necessary.
But also we should likely have a verifier on the dialect itself.

119

Can you replace auto with the types here? (except the first one)

[akuegel]

Address review comments. Start adding more tests.

[akuegel]

Thanks for the review comments :) And sorry for taking so long to respond. I was busy with some other tasks.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1045

I had looked around but didn't find tests for valid constants. I have now added tests for complex constants to Standard/ops.mlir. There were already tests for invalid constants in Standard/invalid.mlir, and I added additional tests to cover the new logic.

mlir/lib/IR/BuiltinAttributes.cpp

589

Thanks for the suggestion :) Done.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

113

So if I understand correctly, most of the checks should be done in the Verifier. Is there another one except LLVM::ConstantOp verifier that I would have to adapt?
I just added an additional check here that the array attribute has size 2. I think the other things are checked via the recursive call.

[mehdi_amini]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

113

Yes you're right, the verifier will be enough.

However the verifier checks an MLIR type:
if (auto complexTy = op.getType().dyn_cast<ComplexType>()) { while here you only have a "struct" as input. So it isn't clear to me that it is totally redundant?

115

What is this catching/exercising by the way?
Why isn't it an error?

[akuegel]

Improve verifier logic, add tests.

[akuegel]

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

335 ↗

(On Diff #344782)

I have added this to be able to return an error for nested structs.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

113

Good point, I didn't realize that I need to check for LLVMStructType in the verifier. I think I understand a bit better now how the interaction of the verifier and ModuleTranslation is. I hope the code I added makes now more sense :)

115

Yes, that should be an error. But this logic is now also part of the verifier.
I have added some logic though for checking whether the created constants that the struct will consist of have the correct type. I can possibly also check this in the verifier instead, but it might be a bit harder to do there.

[akuegel]

Fix assertion error triggered in new test.

[akuegel]

It seems one of the tests I have added triggered an assertion error. I added a check that the FloatAttr has the correct type before calling ConstantFP::get.

Please take another look :)

[akuegel]

Fix ClangTidy warning.

[mehdi_amini]

LG, thanks!

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let description = [{
     ([rationale](../Rationale/Rationale.md#multithreading-the-compiler)).
   }];
 
   let arguments = (ins AnyAttr:$value);
   let results = (outs AnyType);
 
   let builders = [
     OpBuilder<(ins "Attribute":$value),
-    [{ build($_builder, $_state, value.getType(), value); }]>];
+    [{ build($_builder, $_state, value.getType(), value); }]>,
+    OpBuilder<(ins "Attribute":$value, "Type":$type),
+    [{ build($_builder, $_state, type, value); }]>,
+  ];
 
   let extraClassDeclaration = [{
     Attribute getValue() { return (*this)->getAttr("value"); }
 
     /// Returns true if a constant operation can be built with the given value
     /// and result type.
     static bool isBuildableWith(Attribute value, Type type);
   }];

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

if (StringAttr sAttr = op.value().dyn_cast<StringAttr>()) {
     if (!arrayType || arrayType.getNumElements() != sAttr.getValue().size() ||
         !arrayType.getElementType().isInteger(8)) {
       return op->emitOpError()
              << "expected array type of " << sAttr.getValue().size()
              << " i8 elements for the string constant";
     }
     return success();
   }
-  if (!op.value().isa<IntegerAttr, FloatAttr, ElementsAttr>())
+  if (auto arrAttr = op.value().dyn_cast<ArrayAttr>()) {
+    if (arrAttr.size() != 2) {
+      return op.emitOpError() << "expected array with exactly 2 elements";
+    }
+    return success();
+  }
+  if (!op.value().isa<IntegerAttr, ArrayAttr, FloatAttr, ElementsAttr>())
     return op.emitOpError()
            << "only supports integer, float, string or elements attributes";
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Utility functions for parsing atomic ops
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp

for (auto operand : llvm::enumerate(linalgOp.getInputOpOperands())) {
 
       // The operands list is same as the linalgOp with the argument for
       // constant index dropped.
       SmallVector<Value, 4> fusedOperands(linalgOp.getInputs());
       fusedOperands.erase(std::next(fusedOperands.begin(), operand.index()));
 
       // Create a constant scalar value from the splat constant.
       Value scalarConstant = rewriter.create<ConstantOp>(
-          def->getLoc(), constantAttr.getSplatValue());
+          def->getLoc(), constantAttr.getSplatValue(),
+          constantAttr.getType().getElementType());
 
       LinalgOp fusedOp = createLinalgOpOfSameType(
           linalgOp, rewriter, rewriter.getUnknownLoc(),
           linalgOp->getResultTypes(),
           /*inputs=*/fusedOperands,
           /*outputs=*/linalgOp.getOutputs(),
           rewriter.getAffineMapArrayAttr(fusedIndexMaps),
           linalgOp.iterator_types(),

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

 static void print(OpAsmPrinter &p, ConstantOp &op) {
   p << "constant ";
   p.printOptionalAttrDict(op->getAttrs(), /*elidedAttrs=*/{"value"});
 
   if (op->getAttrs().size() > 1)
     p << ' ';
   p << op.getValue();
 
-  // If the value is a symbol reference, print a trailing type.
-  if (op.getValue().isa<SymbolRefAttr>())
+  // If the value is a symbol reference or Array, print a trailing type.
+  if (op.getValue().isa<SymbolRefAttr>() || op.getValue().isa<ArrayAttr>())
     p << " : " << op.getType();
 }
 
 static ParseResult parseConstantOp(OpAsmParser &parser,
                                    OperationState &result) {
   Attribute valueAttr;
   if (parser.parseOptionalAttrDict(result.attributes) ||
       parser.parseAttribute(valueAttr, "value", result.attributes))
     return failure();
 
-  // If the attribute is a symbol reference, then we expect a trailing type.
+  // If the attribute is a symbol reference or array, then we expect a trailing
+  // type.
   Type type;
-  if (!valueAttr.isa<SymbolRefAttr>())
+  if (!valueAttr.isa<SymbolRefAttr>() && !valueAttr.isa<ArrayAttr>())
     type = valueAttr.getType();
   else if (parser.parseColonType(type))
     return failure();
 
   // Add the attribute type to the list.
   return parser.addTypeToList(type, result.types);
 }
 
 /// The constant op requires an attribute, and furthermore requires that it
 /// matches the return type.
 static LogicalResult verify(ConstantOp &op) {
   auto value = op.getValue();
   if (!value)
     return op.emitOpError("requires a 'value' attribute");
 
if (auto intAttr = value.dyn_cast<IntegerAttr>()) {
     unsigned bitwidth = intType.getWidth();
     APInt intVal = intAttr.getValue();
     if (!intVal.isSignedIntN(bitwidth) && !intVal.isIntN(bitwidth))
       return op.emitOpError("requires 'value' to be an integer within the "
                             "range of the integer result type");
     return success();
   }
 
+  if (auto arrayAttr = value.dyn_cast<ArrayAttr>()) {
+    auto complexTy = type.dyn_cast<ComplexType>();
+    if (!complexTy || arrayAttr.size() != 2)
+      return op.emitOpError("requires 'value' to be a complex constant");
+    auto complexEltTy = complexTy.getElementType();
+    if (complexEltTy != arrayAttr[0].getType() ||
+        complexEltTy != arrayAttr[1].getType()) {
+      return op.emitOpError()
+             << "requires attribute's element types (" << arrayAttr[0].getType()
+             << ", " << arrayAttr[1].getType()
+             << ") to match the element type of the op's return type ("
+             << complexEltTy << ")";
+    }
+    return success();
+  }
+
   if (type.isa<FloatType>()) {
     if (!value.isa<FloatAttr>())
       return op.emitOpError("requires 'value' to be a floating point constant");
     return success();
   }
 
   if (type.isa<ShapedType>()) {
     if (!value.isa<ElementsAttr>())

 
 /// Returns true if a constant operation can be built with the given value and
 /// result type.
 bool ConstantOp::isBuildableWith(Attribute value, Type type) {
   // SymbolRefAttr can only be used with a function type.
   if (value.isa<SymbolRefAttr>())
     return type.isa<FunctionType>();
   // The attribute must have the same type as 'type'.
-  if (value.getType() != type)
+  if (!value.getType().isa<NoneType>() && value.getType() != type)
     return false;
   // If the type is an integer type, it must be signless.
   if (IntegerType integerTy = type.dyn_cast<IntegerType>())
     if (!integerTy.isSignless())
       return false;
   // Finally, check that the attribute kind is handled.
-  return value.isa<IntegerAttr, FloatAttr, ElementsAttr, UnitAttr>();
+  return value.isa<IntegerAttr, ArrayAttr, FloatAttr, ElementsAttr, UnitAttr>();
 }
 
 void ConstantFloatOp::build(OpBuilder &builder, OperationState &result,
                             const APFloat &value, FloatType type) {
   ConstantOp::build(builder, result, type, builder.getFloatAttr(type, value));
 }
 
 bool ConstantFloatOp::classof(Operation *op) {

mlir/lib/IR/BuiltinAttributes.cpp

if (auto intEltTy = eltTy.dyn_cast<IntegerType>())
     return IntegerAttr::get(eltTy, *IntElementIterator(owner, index));
   if (eltTy.isa<IndexType>())
     return IntegerAttr::get(eltTy, *IntElementIterator(owner, index));
   if (auto floatEltTy = eltTy.dyn_cast<FloatType>()) {
     IntElementIterator intIt(owner, index);
     FloatElementIterator floatIt(floatEltTy.getFloatSemantics(), intIt);
     return FloatAttr::get(eltTy, *floatIt);
   }
+  if (auto complexTy = eltTy.dyn_cast<ComplexType>()) {
+    auto complexEltTy = complexTy.getElementType();
+    ComplexIntElementIterator complexIntIt(owner, index);
+    if (complexEltTy.isa<IntegerType>()) {
+      auto value = *complexIntIt;
+      auto real = IntegerAttr::get(complexEltTy, value.real());
+      auto imag = IntegerAttr::get(complexEltTy, value.imag());
+      return ArrayAttr::get(complexTy.getContext(),
+                            std::vector<Attribute>{real, imag});
+    }
+
+    ComplexFloatElementIterator complexFloatIt(
+        complexEltTy.cast<FloatType>().getFloatSemantics(), complexIntIt);
+    auto value = *complexFloatIt;
+    auto real = FloatAttr::get(complexEltTy, value.real());
+    auto imag = FloatAttr::get(complexEltTy, value.imag());
+    return ArrayAttr::get(complexTy.getContext(),
+                          std::vector<Attribute>{real, imag});
+  }
   if (owner.isa<DenseStringElementsAttr>()) {
     ArrayRef<StringRef> vals = owner.getRawStringData();
     return StringAttr::get(owner.isSplat() ? vals.front() : vals[index], eltTy);
   }
   llvm_unreachable("unexpected element type");
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

 /// attributes and combinations thereof.  In case of error, report it to `loc`
 /// and return nullptr.
 llvm::Constant *mlir::LLVM::detail::getLLVMConstant(
     llvm::Type *llvmType, Attribute attr, Location loc,
     const ModuleTranslation &moduleTranslation) {
   if (!attr)
     return llvm::UndefValue::get(llvmType);
   if (llvmType->isStructTy()) {
-    emitError(loc, "struct types are not supported in constants");
+    if (auto arrayAttr = attr.dyn_cast<ArrayAttr>()) {
+      auto *structType = cast<::llvm::StructType>(llvmType);
+      auto *elementType = structType->getElementType(0);
+      auto *real =
+          getLLVMConstant(elementType, arrayAttr[0], loc, moduleTranslation);
+      auto *imag =
+          getLLVMConstant(elementType, arrayAttr[1], loc, moduleTranslation);
+      return llvm::ConstantStruct::get(structType, {real, imag});
+    }
+    emitError(loc, "struct types are only supported for complex constants");
     return nullptr;
   }
   // For integer types, we allow a mismatch in sizes as the index type in
   // MLIR might have a different size than the index type in the LLVM module.
   if (auto intAttr = attr.dyn_cast<IntegerAttr>())
     return llvm::ConstantInt::get(
         llvmType,
         intAttr.getValue().sextOrTrunc(llvmType->getIntegerBitWidth()));

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

 // expected-error @+1 {{llvm.align attribute attached to LLVM non-pointer argument}}
 llvm.func @invalid_align(%arg0 : f32 {llvm.align = 4}) -> f32 {
   llvm.return %arg0 : f32
 }
 
 // -----
 
 llvm.func @no_nested_struct() -> !llvm.array<2 x array<2 x array<2 x struct<(i32)>>>> {
-  // expected-error @+1 {{struct types are not supported in constants}}
+  // expected-error @+1 {{struct types are only supported for complex constants}}
   %0 = llvm.mlir.constant(dense<[[[1, 2], [3, 4]], [[42, 43], [44, 45]]]> : tensor<2x2x2xi32>) : !llvm.array<2 x array<2 x array<2 x struct<(i32)>>>>
   llvm.return %0 : !llvm.array<2 x array<2 x array<2 x struct<(i32)>>>>
 }
 
 // -----
 
 // expected-error @+1 {{unsupported constant value}}
 llvm.mlir.global internal constant @test([2.5, 7.4]) : !llvm.array<2 x f64>

mlir/test/Target/LLVMIR/llvmir.mlir

 }
 
 llvm.func @stringconstant() -> !llvm.array<12 x i8> {
   %1 = llvm.mlir.constant("Hello world!") : !llvm.array<12 x i8>
   // CHECK: ret [12 x i8] c"Hello world!"
   llvm.return %1 : !llvm.array<12 x i8>
 }
 
+llvm.func @complexfpconstant() -> !llvm.struct<(f32, f32)> {
+  %1 = llvm.mlir.constant([-1.000000e+00 : f32, 0.000000e+00 : f32]) : !llvm.struct<(f32, f32)>
+  // CHECK: ret { float, float } { float -1.000000e+00, float 0.000000e+00 }
+  llvm.return %1 : !llvm.struct<(f32, f32)>
+}
+
+llvm.func @complexintconstant() -> !llvm.struct<(i32, i32)> {
+  %1 = llvm.mlir.constant([-1 : i32, 0 : i32]) : !llvm.struct<(i32, i32)>
+  // CHECK: ret { i32, i32 } { i32 -1, i32 0 }
+  llvm.return %1 : !llvm.struct<(i32, i32)>
+}
+
 llvm.func @noreach() {
 // CHECK:    unreachable
   llvm.unreachable
 }
 
 // CHECK-LABEL: define void @fcmp
 llvm.func @fcmp(%arg0: f32, %arg1: f32) {
   // CHECK: fcmp oeq float %0, %1