[flang] Add support for dense complex constants

[luporl]

Summary

Add support for representing complex array constants with MLIR
dense attribute. This improves compile time and greatly reduces
memory usage of programs with large complex array constants.

Fixes https://github.com/llvm/llvm-project/issues/63610

[luporl]

Ping

[vzakhari]

flang/test/Lower/array.f90

133

I do not think it is a good idea to mix mlir.complex and fir.complex types here. fir.global might be allowing this right now, but there is no guarantee that mlir.complex and fir.complex will be converted to the same type during fir.global to llvm.global codegen.

Can we at least add an assertion in fir.global codegen that the init attribute data type and the operation's result data type match?

[luporl]

Assert that operation's result type matches init attribute data type.

[luporl]

flang/test/Lower/array.f90

133

Assert added.
There is already a mix between fir.logical and integer types.
For instance, an array constructor with 3 logical elements becomes something like:
fir.global internal @_QQro.3xl4.1(dense<[1, 0, 1]> : tensor<3xi32>) constant : !fir.array<3x!fir.logical<4>>.

[vzakhari]

flang/lib/Optimizer/CodeGen/CodeGen.cpp

2950

I think this assert is a no-op, because tyAttr is the created op's explicit result type. What I suggested to assert is that initAttr's type and tyAttr comply.

[luporl]

flang/lib/Optimizer/CodeGen/CodeGen.cpp

2950

Sorry for the delayed response, I've been away for a few days.

initAttr's type and tyAttr doesn't match.

Considering a complex array of 10 elements, they will be:

tyAttr:   !llvm.array<10 x struct<(f32, f32)>>
initAttr's type: tensor<10xcomplex<f32>>
global: llvm.mlir.global internal @_QFEc1(dense<[(3.000000e+00,4.000000e+00), ...]> : tensor<10xcomplex<f32>>) {addr_space = 0 : i32} : !llvm.array<10 x struct<(f32, f32)>>

It seems the global's dense attribute and its tensor<10xcomplex<f32>> type are lowered to LLVM on a later pass.

It's easy to extract the element type from both (struct<(f32, f32) and complex<f32>), but is there a function that can be called from codegen to convert from MLIR to LLVM type?

[luporl]

Assert that initAttr's type and tyAttr comply.

[luporl]

flang/lib/Optimizer/CodeGen/CodeGen.cpp

2950

A new function was added to compare initAttr's type and tyAttr.

It extracts the element type of both and convert initAttr's element type to its corresponding LLVM type, to make it possible to compare it with tyAttr's type.

The new function is only called by the assert.

flang/test/Lower/array.f90

133

Just expanding a bit more on this, I don't see any good way to avoid mixing mlir.complex and fir.complex types here.

fir.complex can't be changed to mlir.complex without adding casts to the code that manipulates the global's elements, which seems worse. Or that code would need to be changed to operate with mlir.complex values, which would probably require a lot of work.

mlir.complex can't be changed to fir.complex, as it is not a supported attribute type. Maybe it could be wrapped inside an OpaqueAttr, but I don't know if it would work and it doesn't seem a good idea either.

[vzakhari]

Thank you for the extra steps to improve the verification and for the explanations! I think it is as good as it can get currently.

[luporl]

Thank you for the review!

flang/lib/Lower/ConvertConstant.cpp

else {
       static_assert(KIND <= 16, "integers with KIND > 16 are not supported");
       return builder.getIntegerAttr(
           type, llvm::APInt(KIND * 8,
                             {value.ToUInt64(), value.SHIFTR(64).ToUInt64()}));
     }
   } else if constexpr (TC == Fortran::common::TypeCategory::Logical) {
     return builder.getIntegerAttr(type, value.IsTrue());
   } else {
-    static_assert(TC == Fortran::common::TypeCategory::Real,
-                  "type values cannot be converted to attributes");
+    auto getFloatAttr = [&](const auto &value, mlir::Type type) {
       std::string str = value.DumpHexadecimal();
       auto floatVal =
           consAPFloat(builder.getKindMap().getFloatSemantics(KIND), str);
       return builder.getFloatAttr(type, floatVal);
+    };
+
+    if constexpr (TC == Fortran::common::TypeCategory::Real) {
+      return getFloatAttr(value, type);
+    } else {
+      static_assert(TC == Fortran::common::TypeCategory::Complex,
+                    "type values cannot be converted to attributes");
+      mlir::Type eleTy = mlir::cast<mlir::ComplexType>(type).getElementType();
+      llvm::SmallVector<mlir::Attribute, 2> attrs = {
+          getFloatAttr(value.REAL(), eleTy),
+          getFloatAttr(value.AIMAG(), eleTy)};
+      return builder.getArrayAttr(attrs);
+    }
   }
   return {};
 }
 
 namespace {
 /// Helper class to lower an array constant to a global with an MLIR dense
 /// attribute.
 ///
-/// If we have an array of integer, real, or logical, then we can
+/// If we have an array of integer, real, complex, or logical, then we can
 /// create a global array with the dense attribute.
 ///
-/// The mlir tensor type can only handle integer, real, or logical. It
-/// does not currently support nested structures which is required for
-/// complex.
+/// The mlir tensor type can only handle integer, real, complex, or logical.
+/// It does not currently support nested structures.
 class DenseGlobalBuilder {
 public:
   static fir::GlobalOp tryCreating(fir::FirOpBuilder &builder,
                                    mlir::Location loc, mlir::Type symTy,
                                    llvm::StringRef globalName,
                                    mlir::StringAttr linkage, bool isConst,
                                    const Fortran::lower::SomeExpr &initExpr) {
     DenseGlobalBuilder globalBuilder;
     std::visit(
         Fortran::common::visitors{
             [&](const Fortran::evaluate::Expr<Fortran::evaluate::SomeLogical> &
                     x) { globalBuilder.tryConvertingToAttributes(builder, x); },
             [&](const Fortran::evaluate::Expr<Fortran::evaluate::SomeInteger> &
                     x) { globalBuilder.tryConvertingToAttributes(builder, x); },
             [&](const Fortran::evaluate::Expr<Fortran::evaluate::SomeReal> &x) {
               globalBuilder.tryConvertingToAttributes(builder, x);
             },
+            [&](const Fortran::evaluate::Expr<Fortran::evaluate::SomeComplex> &
+                    x) { globalBuilder.tryConvertingToAttributes(builder, x); },
             [](const auto &) {},
         },
         initExpr.u);
     return globalBuilder.tryCreatingGlobal(builder, loc, symTy, globalName,
                                            linkage, isConst);
   }
 
   template <Fortran::common::TypeCategory TC, int KIND>
void tryConvertingToAttributes(
           &constant) {
     static_assert(TC != Fortran::common::TypeCategory::Character,
                   "must be numerical or logical");
     auto attrTc = TC == Fortran::common::TypeCategory::Logical
                       ? Fortran::common::TypeCategory::Integer
                       : TC;
     attributeElementType = Fortran::lower::getFIRType(
         builder.getContext(), attrTc, KIND, std::nullopt);
+    if (auto firCTy = mlir::dyn_cast<fir::ComplexType>(attributeElementType))
+      attributeElementType =
+          mlir::ComplexType::get(firCTy.getEleType(builder.getKindMap()));
     for (auto element : constant.values())
       attributes.push_back(
           convertToAttribute<TC, KIND>(builder, element, attributeElementType));
   }
 
   /// Try converting an evaluate::Expr to a list of MLIR attributes.
   template <typename SomeCat>
   void tryConvertingToAttributes(fir::FirOpBuilder &builder,
llvm::StringRef globalName = converter.getUniqueLitName(
       eleTy);
   fir::GlobalOp global = builder.getNamedGlobal(globalName);
   if (!global) {
     // Using a dense attribute for the initial value instead of creating an
     // intialization body speeds up MLIR/LLVM compilation, but this is not
     // always possible.
     if constexpr (T::category == Fortran::common::TypeCategory::Logical ||
                   T::category == Fortran::common::TypeCategory::Integer ||
-                  T::category == Fortran::common::TypeCategory::Real) {
+                  T::category == Fortran::common::TypeCategory::Real ||
+                  T::category == Fortran::common::TypeCategory::Complex) {
       global = DenseGlobalBuilder::tryCreating(
           builder, loc, arrayTy, globalName, builder.createInternalLinkage(),
           true, constant);
     }
     if (!global)
       // If the number of elements of the array is huge, the compilation may
       // use a lot of memory and take a very long time to complete.
       // Empirical evidence shows that an array with 150000 elements of

flang/lib/Lower/ConvertVariable.cpp

static fir::GlobalOp defineGlobal(Fortran::lower::AbstractConverter &converter,
   if (global && globalIsInitialized(global))
     return global;
 
   if (Fortran::semantics::IsProcedurePointer(sym))
     TODO(loc, "procedure pointer globals");
 
   // If this is an array, check to see if we can use a dense attribute
   // with a tensor mlir type. This optimization currently only supports
-  // Fortran arrays of integer, real, or logical. The tensor type does
-  // not support nested structures which are needed for complex numbers.
+  // Fortran arrays of integer, real, complex, or logical. The tensor
+  // type does not support nested structures.
   if (symTy.isa<fir::SequenceType>() &&
       !Fortran::semantics::IsAllocatableOrPointer(sym)) {
     mlir::Type eleTy = symTy.cast<fir::SequenceType>().getEleTy();
-    if (eleTy.isa<mlir::IntegerType, mlir::FloatType, fir::LogicalType>()) {
+    if (eleTy.isa<mlir::IntegerType, mlir::FloatType, fir::ComplexType,
+                  fir::LogicalType>()) {
       const auto *details =
           sym.detailsIf<Fortran::semantics::ObjectEntityDetails>();
       if (details->init()) {
         global = Fortran::lower::tryCreatingDenseGlobal(
             builder, loc, symTy, globalName, linkage, isConst,
             details->init().value());
         if (global) {
           global.setVisibility(mlir::SymbolTable::Visibility::Public);

flang/lib/Optimizer/CodeGen/CodeGen.cpp

struct HasValueOpConversion : public FIROpConversion<fir::HasValueOp> {
   matchAndRewrite(fir::HasValueOp op, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<mlir::LLVM::ReturnOp>(op,
                                                       adaptor.getOperands());
     return mlir::success();
   }
 };
 
+// Check if attr's type is compatible with ty.
+//
+// This is done by comparing attr's element type, converted to LLVM type,
+// with ty's element type.
+//
+// Only integer and floating point (including complex) attributes are
+// supported. Also, attr is expected to have a TensorType and ty is expected
+// to be of LLVMArrayType. If any of the previous conditions is false, then
+// the specified attr and ty are not supported by this function and are
+// assumed to be compatible.
+static inline bool attributeTypeIsCompatible(mlir::MLIRContext *ctx,
+                                             mlir::Attribute attr,
+                                             mlir::Type ty) {
+  // Get attr's LLVM element type.
+  if (!attr)
+    return true;
+  auto intOrFpEleAttr = mlir::dyn_cast<mlir::DenseIntOrFPElementsAttr>(attr);
+  if (!intOrFpEleAttr)
+    return true;
+  auto tensorTy = mlir::dyn_cast<mlir::TensorType>(intOrFpEleAttr.getType());
+  if (!tensorTy)
+    return true;
+  mlir::Type attrEleTy =
+      mlir::LLVMTypeConverter(ctx).convertType(tensorTy.getElementType());
+
+  // Get ty's element type.
+  auto arrTy = mlir::dyn_cast<mlir::LLVM::LLVMArrayType>(ty);
+  if (!arrTy)
+    return true;
+  mlir::Type eleTy = arrTy.getElementType();
+  while ((arrTy = mlir::dyn_cast<mlir::LLVM::LLVMArrayType>(eleTy)))
+    eleTy = arrTy.getElementType();
+
+  return attrEleTy == eleTy;
+}
+
 /// Lower `fir.global` operation to `llvm.global` operation.
 /// `fir.insert_on_range` operations are replaced with constant dense attribute
 /// if they are applied on the full range.
 struct GlobalOpConversion : public FIROpConversion<fir::GlobalOp> {
   using FIROpConversion::FIROpConversion;
 
   mlir::LogicalResult
   matchAndRewrite(fir::GlobalOp global, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     auto tyAttr = convertType(global.getType());
     if (global.getType().isa<fir::BaseBoxType>())
       tyAttr = tyAttr.cast<mlir::LLVM::LLVMPointerType>().getElementType();
     auto loc = global.getLoc();
     mlir::Attribute initAttr = global.getInitVal().value_or(mlir::Attribute());
+    assert(attributeTypeIsCompatible(global.getContext(), initAttr, tyAttr));
     auto linkage = convertLinkage(global.getLinkName());
     auto isConst = global.getConstant().has_value();
     auto g = rewriter.create<mlir::LLVM::GlobalOp>(
         loc, tyAttr, isConst, linkage, global.getSymName(), initAttr);
 
     auto module = global->getParentOfType<mlir::ModuleOp>();
     // Add comdat if necessary
     if (fir::getTargetTriple(module).supportsCOMDAT() &&
         (linkage == mlir::LLVM::Linkage::Linkonce ||
          linkage == mlir::LLVM::Linkage::LinkonceODR)) {
       addComdat(g, rewriter, module);
     }
 

flang/test/Lower/array.f90

 
 ! CHECK-LABEL range
 subroutine range()
   ! Compile-time initalized arrays
   integer, dimension(10) :: a0
   real, dimension(2,3) ::  a1
   integer, dimension(3,4) :: a2
   integer, dimension(2,3,4) :: a3
- 
+  complex, dimension(2,3) :: c0, c1
   a0 = (/1, 2, 3, 3, 3, 3, 3, 3, 3, 3/)
   a1 = reshape((/3.5, 3.5, 3.5, 3.5, 3.5, 3.5/), shape(a1))
   a2 = reshape((/1, 3, 3, 5, 3, 3, 3, 3, 9, 9, 9, 8/), shape(a2))
   a3 = reshape((/1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12/), shape(a3))
+
+  c0 = reshape((/(1.0, 1.5), (2.0, 2.5), (3.0, 3.5), (4.0, 4.5), (5.0, 5.5), (6.0, 6.5)/), shape(c0))
+  data c1/6 * (0.0, 0.0)/
 end subroutine range
 
+! c1 data
+! CHECK: fir.global internal @_QFrangeEc1(dense<(0.000000e+00,0.000000e+00)> : tensor<3x2xcomplex<f32>>) : !fir.array<2x3x!fir.complex<4>>
+
 ! a0 array constructor
 ! CHECK: fir.global internal @_QQro.10xi4.{{.*}}(dense<[1, 2, 3, 3, 3, 3, 3, 3, 3, 3]> : tensor<10xi32>) constant : !fir.array<10xi32>
 
 ! a1 array constructor
 ! CHECK: fir.global internal @_QQro.2x3xr4.{{.*}}(dense<3.500000e+00> : tensor<3x2xf32>) constant : !fir.array<2x3xf32>
 
 ! a2 array constructor
 ! CHECK: fir.global internal @_QQro.3x4xi4.{{.*}}(dense<{{\[\[1, 3, 3], \[5, 3, 3], \[3, 3, 9], \[9, 9, 8]]}}> : tensor<4x3xi32>) constant : !fir.array<3x4xi32>
 
 ! a3 array constructor
 ! CHECK: fir.global internal @_QQro.2x3x4xi4.{{.*}}(dense<{{\[\[\[1, 1], \[2, 2], \[3, 3]], \[\[4, 4], \[5, 5], \[6, 6]], \[\[7, 7], \[8, 8], \[9, 9]], \[\[10, 10], \[11, 11], \[12, 12]]]}}> : tensor<4x3x2xi32>) constant : !fir.array<2x3x4xi32>
 
+! c0 array constructor
+! CHECK: fir.global internal @_QQro.2x3xz4.{{.*}}(dense<{{\[}}[(1.000000e+00,1.500000e+00), (2.000000e+00,2.500000e+00)], [(3.000000e+00,3.500000e+00), (4.000000e+00,4.500000e+00)], [(5.000000e+00,5.500000e+00), (6.000000e+00,6.500000e+00)]]> : tensor<3x2xcomplex<f32>>) constant : !fir.array<2x3x!fir.complex<4>>
+
 ! CHECK-LABEL rangeGlobal
 subroutine rangeGlobal()
 ! CHECK: fir.global internal @_QFrangeglobal{{.*}}(dense<[1, 1, 2, 2, 3, 3]> : tensor<6xi32>) : !fir.array<6xi32>
   integer, dimension(6) :: a0 = (/ 1, 1, 2, 2, 3, 3 /)
 
 end subroutine rangeGlobal
 
 ! CHECK-LABEL hugeGlobal