[flang] Handle array constants of any rank

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

Reviewed By: jeanPerier

Differential Revision: https://reviews.llvm.org/D150686

diff --git a/flang/lib/Lower/ConvertConstant.cpp b/flang/lib/Lower/ConvertConstant.cpp
index 7afcc96..de8830a 100644 (file)
--- a/flang/lib/Lower/ConvertConstant.cpp
+++ b/flang/lib/Lower/ConvertConstant.cpp
@@ -20,6 +20,8 @@
 #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,
@@ -66,17 +68,12 @@ 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,
@@ -124,8 +121,6 @@ private:
           &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;
@@ -158,12 +153,16 @@ private:
                                   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);
   }
@@ -544,6 +543,13 @@ genOutlineArrayLit(Fortran::lower::AbstractConverter &converter,
           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) {

diff --git a/flang/lib/Lower/ConvertVariable.cpp b/flang/lib/Lower/ConvertVariable.cpp
index 2204505..b0dcf4d 100644 (file)
--- a/flang/lib/Lower/ConvertVariable.cpp
+++ b/flang/lib/Lower/ConvertVariable.cpp
@@ -431,14 +431,9 @@ static fir::GlobalOp defineGlobal(Fortran::lower::AbstractConverter &converter,
 
   // 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>()) {

diff --git a/flang/test/Lower/array.f90 b/flang/test/Lower/array.f90
index 9d15b3b..c64da68 100644 (file)
--- a/flang/test/Lower/array.f90
+++ b/flang/test/Lower/array.f90
@@ -102,33 +102,25 @@ subroutine range()
   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()
@@ -137,6 +129,15 @@ subroutine rangeGlobal()
 
 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

diff --git a/flang/test/Lower/dense-array-any-rank.f90 b/flang/test/Lower/dense-array-any-rank.f90
new file mode 100644 (file)
index 0000000..437fdec
--- /dev/null
+++ b/flang/test/Lower/dense-array-any-rank.f90
@@ -0,0 +1,25 @@
+! 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]]]}}