[flang] Handle array constants of any rank

[luporl]

Summary

Add support for representing array constants of any rank with MLIR
dense attribute. This greatly improves compile time and memory
usage of programs with large array constants. We still support only
arrays of a few basic types, such as integer, real and logic.

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

[clementval]

Thanks for working on this. Did you check with the example given by @jeanPerier in the issue? What kind of build time to you see?

[luporl]

In D150686#4346701, @clementval wrote:

Thanks for working on this. Did you check with the example given by @jeanPerier in the issue? What kind of build time to you see?

Yes, I ended up commenting more about this in the issue. With @jeanPerier's example, I get the following build time:

$ time flang -c ftest_reduced.f90

real	0m1.059s
user	0m0.962s
sys	0m0.096s

This is on an ARMv8.2 Ampere Altra Mt. Jade machine.

(In the command above, flang is actually an alias to /home/leandro.lupori/git/flang-i60376/buildr/bin/flang-new)

[clementval]

It might to be good to add a test with the codegen to make sure we have what we expect at the LLVM IR level.

[jeanPerier]

Thank you very much for improving this limitation!

flang/lib/Lower/ConvertConstant.cpp

553

I know I initially suggested the TODO, but after discussions, it would be too harsh to artificially limit the size of a program that can be compile. Some people with more RAM and time may just be OK with bigger compile time here to be able to test/use flang without modifying their programs.

The ideal would be to add a driver switch to control the size limit and to tell the user to use that switch to increase the limit in the TODO message that would warn this may lead to very long compile time.

Since threading a switch here may not be straightforward, I am also OK with just removing the TODO for now so that flang can be tested with no limit and to leave your comment with the addition that it would be nice to add a switch to control an array size after which flang should not continue to compile.

Besides, after your patch it should be a lot less common to hit the compilation time/memory issue (most big array constant expressions will use integer and reals out there).

[luporl]

Add test that checks LLVM IR matches what is expected

[luporl]

flang/lib/Lower/ConvertConstant.cpp

553

Ok, since big array constants of types other than integer and real shouldn't be common, for now I'll just remove the TODO and leave a comment.

[luporl]

Remove TODO and suggest the addition of a switch to control the maximum array size supported.

[jeanPerier]

Thanks a lot, LGTM

[luporl]

Fix dense attributes for integers with KIND > 8

[jeanPerier]

Thanks for the fix

flang/lib/Lower/ConvertConstant.cpp

 #include "flang/Lower/AbstractConverter.h"
 #include "flang/Lower/BuiltinModules.h"
 #include "flang/Lower/ConvertType.h"
 #include "flang/Lower/ConvertVariable.h"
 #include "flang/Lower/Mangler.h"
 #include "flang/Optimizer/Builder/Complex.h"
 #include "flang/Optimizer/Builder/Todo.h"
 
+#include <algorithm>
+
 /// Convert string, \p s, to an APFloat value. Recognize and handle Inf and
 /// NaN strings as well. \p s is assumed to not contain any spaces.
 static llvm::APFloat consAPFloat(const llvm::fltSemantics &fsem,
                                  llvm::StringRef s) {
   assert(!s.contains(' '));
   if (s.compare_insensitive("-inf") == 0)
     return llvm::APFloat::getInf(fsem, /*negative=*/true);
   if (s.compare_insensitive("inf") == 0 || s.compare_insensitive("+inf") == 0)
static mlir::Attribute convertToAttribute(
   }
   return {};
 }
 
 namespace {
 /// Helper class to lower an array constant to a global with an MLIR dense
 /// attribute.
 ///
-/// If we have a rank-1 array of integer, real, or logical, then we can
+/// If we have an array of integer, real, 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.
-///
-/// Also, we currently handle just rank-1 since tensor type assumes
-/// row major array ordering. We will need to reorder the dimensions
-/// in the tensor type to support Fortran's column major array ordering.
-/// How to create this tensor type is to be determined.
 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;
private:
   /// Try converting an evaluate::Constant to a list of MLIR attributes.
   template <Fortran::common::TypeCategory TC, int KIND>
   void tryConvertingToAttributes(
       fir::FirOpBuilder &builder,
       const Fortran::evaluate::Constant<Fortran::evaluate::Type<TC, KIND>>
           &constant) {
     static_assert(TC != Fortran::common::TypeCategory::Character,
                   "must be numerical or logical");
-    if (constant.Rank() != 1)
-      return;
     auto attrTc = TC == Fortran::common::TypeCategory::Logical
                       ? Fortran::common::TypeCategory::Integer
                       : TC;
     attributeElementType = Fortran::lower::getFIRType(
         builder.getContext(), attrTc, KIND, std::nullopt);
     for (auto element : constant.values())
       attributes.push_back(
           convertToAttribute<TC, KIND>(builder, element, attributeElementType));
private:
 
   /// Create a fir::Global if MLIR attributes have been successfully created by
   /// tryConvertingToAttributes.
   fir::GlobalOp tryCreatingGlobal(fir::FirOpBuilder &builder,
                                   mlir::Location loc, mlir::Type symTy,
                                   llvm::StringRef globalName,
                                   mlir::StringAttr linkage,
                                   bool isConst) const {
-    // Not a rank 1 "trivial" intrinsic constant array, or empty array.
+    // Not a "trivial" intrinsic constant array, or empty array.
     if (!attributeElementType || attributes.empty())
       return {};
 
+    assert(symTy.isa<fir::SequenceType>() && "expecting an array global");
+    auto arrTy = symTy.cast<fir::SequenceType>();
+    llvm::SmallVector<int64_t> tensorShape(arrTy.getShape());
+    std::reverse(tensorShape.begin(), tensorShape.end());
     auto tensorTy =
-        mlir::RankedTensorType::get(attributes.size(), attributeElementType);
+        mlir::RankedTensorType::get(tensorShape, attributeElementType);
     auto init = mlir::DenseElementsAttr::get(tensorTy, attributes);
     return builder.createGlobal(loc, symTy, globalName, linkage, init, isConst);
   }
 
   llvm::SmallVector<mlir::Attribute> attributes;
   mlir::Type attributeElementType;
 };
 } // namespace
if (!global) {
     if constexpr (T::category == Fortran::common::TypeCategory::Logical ||
                   T::category == Fortran::common::TypeCategory::Integer ||
                   T::category == Fortran::common::TypeCategory::Real) {
       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
+      // complex type takes roughly 30 seconds to compile and uses 4GB of RAM,
+      // on a modern machine.
+      // It would be nice to add a driver switch to control the array size
+      // after which flang should not continue to compile.
       global = builder.createGlobalConstant(
           loc, arrayTy, globalName,
           [&](fir::FirOpBuilder &builder) {
             mlir::Value result =
                 genInlinedArrayLit(converter, loc, arrayTy, constant);
             builder.create<fir::HasValueOp>(loc, result);
           },
           builder.createInternalLinkage());
   }

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
-  // rank-1 Fortran arrays of integer, real, or logical. The tensor
-  // type does not support nested structures which are needed for
-  // complex numbers.
-  // To get multidimensional arrays to work, we will have to use column major
-  // array ordering with the tensor type (so it matches column major ordering
-  // with the Fortran fir.array). By default, tensor types assume row major
-  // ordering. How to create this tensor type is to be determined.
-  if (symTy.isa<fir::SequenceType>() && sym.Rank() == 1 &&
+  // Fortran arrays of integer, real, or logical. The tensor type does
+  // not support nested structures which are needed for complex numbers.
+  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>()) {
       const auto *details =
           sym.detailsIf<Fortran::semantics::ObjectEntityDetails>();
       if (details->init()) {
         global = Fortran::lower::tryCreatingDenseGlobal(
             builder, loc, symTy, globalName, linkage, isConst,

flang/test/Lower/array.f90

 end subroutine s
 
 ! 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
  
   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))
 end subroutine range
 
 ! 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.{{.*}} constant : !fir.array<2x3xf32> {
-  ! CHECK-DAG: %cst = arith.constant {{.*}} : f32
-  ! CHECK: %{{.*}} = fir.insert_on_range %{{[0-9]+}}, %cst from (0, 0) to (1, 2) :
+! 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.{{.*}} constant : !fir.array<3x4xi32> {
-  ! CHECK-DAG: %[[c1_i32:.*]] = arith.constant 1 : i32
-  ! CHECK-DAG: %[[c3_i32:.*]] = arith.constant 3 : i32
-  ! CHECK-DAG: %[[c5_i32:.*]] = arith.constant 5 : i32
-  ! CHECK-DAG: %[[c8_i32:.*]] = arith.constant 8 : i32
-  ! CHECK-DAG: %[[c9_i32:.*]] = arith.constant 9 : i32
-  ! CHECK: %[[r1:.*]] = fir.insert_value %{{.*}}, %{{.*}}, [0 : index, 0 : index] :
-  ! CHECK: %[[r2:.*]] = fir.insert_on_range %[[r1]], %[[c3_i32]] from (1, 0) to (2, 0) :
-  ! CHECK: %[[r3:.*]] = fir.insert_value %[[r2]], %{{.*}}, [0 : index, 1 : index] :
-  ! CHECK: %[[r4:.*]] = fir.insert_on_range %[[r3]], %[[c3_i32]] from (1, 1) to (1, 2) :
-  ! CHECK: %[[r5:.*]] = fir.insert_on_range %[[r4]], %[[c9_i32]] from (2, 2) to (1, 3) :
-  ! CHECK: %[[r6:.*]] = fir.insert_value %[[r5]], %{{.*}}, [2 : index, 3 : index] :
+! 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>
 
 ! 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
+subroutine hugeGlobal()
+  integer, parameter :: D = 500
+  integer, dimension(D, D) :: a
+
+! CHECK: fir.global internal @_QQro.500x500xi4.{{.*}}(dense<{{.*}}> : tensor<500x500xi32>) constant : !fir.array<500x500xi32>
+  a = reshape((/(i, i = 1, D * D)/), shape(a))
+end subroutine hugeGlobal
+
 block data
   real(selected_real_kind(6)) :: x(5,5)
   common /block/ x
   data x(1,1), x(2,1), x(3,1) / 1, 1, 0 /
   data x(1,2), x(2,2), x(4,2) / 1, 1, 2.4 /
   data x(1,3), x(2,3), x(4,3) / 1, 1, 2.4 /
   data x(4,4) / 2.4 /
 end

flang/test/Lower/dense-array-any-rank.f90

+! RUN: bbc -emit-fir -o - %s | FileCheck --check-prefixes="CHECK-FIR" %s
+! RUN: %flang_fc1 -emit-llvm -o - %s | FileCheck --check-prefixes="CHECK-LLVMIR" %s
+
+! CHECK-LABEL: test
+subroutine test()
+  integer, dimension(10) :: a1
+  integer, dimension(3,4) :: a2
+  integer, dimension(2,3,4) :: a3
+
+  a1 = (/1, 2, 3, 4, 5, 6, 7, 8, 9, 10/)
+  a2 = reshape((/11, 12, 13, 21, 22, 23, 31, 32, 33, 41, 42, 43/), shape(a2))
+  a3 = reshape((/111, 112, 121, 122, 131, 132, 211, 212, 221, 222, 231, 232, 311, 312, 321, 322, 331, 332, 411, 412, 421, 422, 431, 432/), shape(a3))
+end subroutine
+
+! a1 array constructor
+! CHECK-FIR: fir.global internal @_QQro.10xi4.{{.*}}(dense<[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]> : tensor<10xi32>) constant : !fir.array<10xi32>
+! CHECK-LLVMIR: @_QQro.10xi4.0 = internal constant [10 x i32] [i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8, i32 9, i32 10]
+
+! a2 array constructor
+! CHECK-FIR: fir.global internal @_QQro.3x4xi4.{{.*}}(dense<{{\[\[11, 12, 13], \[21, 22, 23], \[31, 32, 33], \[41, 42, 43]]}}> : tensor<4x3xi32>) constant : !fir.array<3x4xi32>
+! CHECK-LLVMIR: @_QQro.3x4xi4.1 = internal constant [4 x [3 x i32]] {{\[\[3 x i32] \[i32 11, i32 12, i32 13], \[3 x i32] \[i32 21, i32 22, i32 23], \[3 x i32] \[i32 31, i32 32, i32 33], \[3 x i32] \[i32 41, i32 42, i32 43]]}}
+
+! a3 array constructor
+! CHECK-FIR: fir.global internal @_QQro.2x3x4xi4.{{.*}}(dense<{{\[\[\[111, 112], \[121, 122], \[131, 132]], \[\[211, 212], \[221, 222], \[231, 232]], \[\[311, 312], \[321, 322], \[331, 332]], \[\[411, 412], \[421, 422], \[431, 432]]]}}> : tensor<4x3x2xi32>) constant : !fir.array<2x3x4xi32>
+! CHECK-LLVMIR: @_QQro.2x3x4xi4.2 = internal constant [4 x [3 x [2 x i32]]] {{\[\[3 x \[2 x i32]] \[\[2 x i32] \[i32 111, i32 112], \[2 x i32] \[i32 121, i32 122], \[2 x i32] \[i32 131, i32 132]], \[3 x \[2 x i32]] \[\[2 x i32] \[i32 211, i32 212], \[2 x i32] \[i32 221, i32 222], \[2 x i32] \[i32 231, i32 232]], \[3 x \[2 x i32]] \[\[2 x i32] \[i32 311, i32 312], \[2 x i32] \[i32 321, i32 322], \[2 x i32] \[i32 331, i32 332]], \[3 x \[2 x i32]] \[\[2 x i32] \[i32 411, i32 412], \[2 x i32] \[i32 421, i32 422], \[2 x i32] \[i32 431, i32 432]]]}}