[flang] Fix overallocation by fir-to-llvm-ir pass

[rovka]

Summary

When converting a fir.alloca of an array to the LLVM dialect, we used to
multiply the allocated size by all the constant factors encoded in the
array type. This is fine when the array type is converted to the element
type for the purposes of the allocation, but if it's converted to an
array type, then we might be allocating too much space. For example, for
%2 = fir.alloca !fir.array<8x16x32xf32>, %0, %1 we would allocate
%0 * %1 * 8 * 16 * 32 x llvm.array<32 x array<16 * array<8 x f32>>>. We
really only need to allocate %0 * %1 such arrays.

This patch fixes the issue by taking note of the array type that we're
trying to allocate. It tries to match the behaviour of
LLVMTypeConverter::convertPointerLike, which returns a pointer to the
element type when the array type doesn't have a constant interior.
We consequently only multiply with the constant factors in the array
type if the array type doesn't have a constant interior.

This has the nice side effect that it gets rid of some redundant
multiplications with the constant 1 in some cases.

[schweitz]

flang/lib/Optimizer/CodeGen/CodeGen.cpp

376–384

This does not appear to be the correct.

alloca !fir.array<*:i32>, %1, %2

has a constant sized interior (the element type i32 has a constant size), but it is clearly not valid to drop the array shape arguments.

This applies to a case like

alloca !fir.array<5x6x?x?xf64>, %3, %4, %5

as well. In this second case, space of %3 * %4 * %5 * 5 * 6 * sizeof(f64) bytes are required to be allocated even though the type has a constant interior of 5 * 6 * sizeof(f64).

[rovka]

flang/lib/Optimizer/CodeGen/CodeGen.cpp

376–384

If I understand correctly, the first example is not valid, since you can't call alloca on a sequence type with unknown rank. You'd get an error from here.

As for the second example, it's very similar to the test I'm adding at line 1119, alloca_const_interior_array. That has

fir.alloca !fir.array<8x9x?x?xf32>, %0, %1

and it allocates %0 * %1 * array<9 x array<8 x f32>>, so I believe it's allocating the right amount. Am I missing something?

[schweitz]

flang/lib/Optimizer/CodeGen/CodeGen.cpp

376–384

Ok, got it now. Thanks.

[awarzynski]

Makes sense, thank you @rovka !

flang/lib/Optimizer/CodeGen/CodeGen.cpp

if (alloc.hasLenParams()) {
       } else {
         return emitError(loc, "unexpected type ")
                << scalarType << " with type parameters";
       }
     }
     if (alloc.hasShapeOperands()) {
       mlir::Type allocEleTy = fir::unwrapRefType(alloc.getType());
       // Scale the size by constant factors encoded in the array type.
+      // We only do this for arrays that don't have a constant interior, since
+      // those are the only ones that get decayed to a pointer to the element
+      // type.
       if (auto seqTy = allocEleTy.dyn_cast<fir::SequenceType>()) {
+        if (!seqTy.hasConstantInterior()) {
           fir::SequenceType::Extent constSize = 1;
           for (auto extent : seqTy.getShape())
             if (extent != fir::SequenceType::getUnknownExtent())
               constSize *= extent;
           mlir::Value constVal{
               genConstantIndex(loc, ity, rewriter, constSize).getResult()};
           size = rewriter.create<mlir::LLVM::MulOp>(loc, ity, size, constVal);
         }
+      }
       unsigned end = operands.size();
       for (; i < end; ++i)
         size = rewriter.create<mlir::LLVM::MulOp>(
             loc, ity, size, integerCast(loc, rewriter, ity, operands[i]));
     }
     if (ty == resultTy) {
       // Do not emit the bitcast if ty and resultTy are the same.
       rewriter.replaceOpWithNewOp<mlir::LLVM::AllocaOp>(alloc, ty, size,

flang/test/Fir/convert-to-llvm.fir

 
 // CHECK-LABEL: llvm.func @alloca_ptr_to_array() -> !llvm.ptr<ptr<i32>>
 // CHECK: [[ONE:%.*]] = llvm.mlir.constant(1 : i64) : i64
 // CHECK: [[A:%.*]] = llvm.alloca [[ONE]] x !llvm.ptr<i32>
 // CHECK: llvm.return [[A]] : !llvm.ptr<ptr<i32>>
 
 // -----
 
-// Test fir.alloca of char array
+// Test fir.alloca of array of unknown-length chars
 
 func @alloca_char_array(%l: i32, %e : index) -> !fir.ref<!fir.array<?x?x!fir.char<1,?>>> {
   %a = fir.alloca !fir.array<?x?x!fir.char<1,?>>(%l : i32), %e, %e
   return %a :  !fir.ref<!fir.array<?x?x!fir.char<1,?>>>
 }
 
 // CHECK-LABEL: llvm.func @alloca_char_array
 // CHECK-SAME: ([[L:%.*]]: i32, [[E:%.*]]: i64) -> !llvm.ptr<i8>
 // CHECK-DAG: [[UNUSEDONE:%.*]] = llvm.mlir.constant(1 : i64) : i64
 // CHECK-DAG: [[LCAST:%.*]] = llvm.sext [[L]] : i32 to i64
-// CHECK-DAG: [[ONE:%.*]] = llvm.mlir.constant(1 : i64) : i64
-// CHECK: [[PROD1:%.*]] = llvm.mul [[LCAST]], [[ONE]] : i64
+// CHECK: [[PROD1:%.*]] = llvm.mul [[LCAST]], [[E]] : i64
 // CHECK: [[PROD2:%.*]] = llvm.mul [[PROD1]], [[E]] : i64
-// CHECK: [[PROD3:%.*]] = llvm.mul [[PROD2]], [[E]] : i64
-// CHECK: [[A:%.*]] = llvm.alloca [[PROD3]] x i8 {in_type = !fir.array<?x?x!fir.char<1,?>>
+// CHECK: [[A:%.*]] = llvm.alloca [[PROD2]] x i8 {in_type = !fir.array<?x?x!fir.char<1,?>>
 // CHECK: return [[A]] : !llvm.ptr<i8>
 
 // -----
 
+// Test fir.alloca of array of known-length chars
+
+func @alloca_fixed_char_array(%e : index) -> !fir.ref<!fir.array<?x?x!fir.char<1,8>>> {
+  %a = fir.alloca !fir.array<?x?x!fir.char<1,8>>, %e, %e
+  return %a :  !fir.ref<!fir.array<?x?x!fir.char<1,8>>>
+}
+
+// CHECK-LABEL: llvm.func @alloca_fixed_char_array
+// CHECK-SAME: ([[E:%.*]]: i64) -> !llvm.ptr<array<8 x i8>>
+// CHECK-DAG: [[ONE:%.*]] = llvm.mlir.constant(1 : i64) : i64
+// CHECK: [[PROD1:%.*]] = llvm.mul [[ONE]], [[E]] : i64
+// CHECK: [[PROD2:%.*]] = llvm.mul [[PROD1]], [[E]] : i64
+// CHECK: [[A:%.*]] = llvm.alloca [[PROD2]] x !llvm.array<8 x i8> {in_type = !fir.array<?x?x!fir.char<1,8>>
+// CHECK: return [[A]] : !llvm.ptr<array<8 x i8>>
+
+// -----
+
 // Test fir.alloca of record type with LEN parameters
 //   type t(p1,p2)
 //      integer, len :: p1
 //      integer(kind=2), len :: p2
 //      integer f1
 //      real f2
 //   end type t
 
func @alloca_multidim_array(%0 : index) -> !fir.ref<!fir.array<8x16x32xf32>> {
   %2 = fir.alloca !fir.array<8x16x32xf32>, %0, %1
   return %2 : !fir.ref<!fir.array<8x16x32xf32>>
 }
 
 // CHECK-LABEL: llvm.func @alloca_multidim_array
 // CHECK-SAME: ([[OP1:%.*]]: i64) -> !llvm.ptr<array<32 x array<16 x array<8 x f32>
 // CHECK: [[OP2:%.*]] = llvm.mlir.constant(24 : index) : i64
 // CHECK: [[ONE:%.*]] = llvm.mlir.constant(1 : i64) : i64
-// CHECK: [[ALL:%.*]] = llvm.mlir.constant(4096 : i64) : i64
-// CHECK: [[MUL1:%.*]] = llvm.mul [[ONE]], [[ALL]] : i64
-// CHECK: [[MUL2:%.*]] = llvm.mul [[MUL1]], [[OP1]] : i64
-// CHECK: [[TOTAL:%.*]] = llvm.mul [[MUL2]], [[OP2]] : i64
+// CHECK: [[MUL1:%.*]] = llvm.mul [[ONE]], [[OP1]] : i64
+// CHECK: [[TOTAL:%.*]] = llvm.mul [[MUL1]], [[OP2]] : i64
 // CHECK: [[A:%.*]] = llvm.alloca [[TOTAL]] x !llvm.array<32 x array<16 x array<8 x f32>
 // CHECK: llvm.return [[A]] : !llvm.ptr<array<32 x array<16 x array<8 x f32>
 
 // -----
 
+// Test fir.alloca of a multidimensional array with constant interior
+
+func @alloca_const_interior_array(%0 : index) -> !fir.ref<!fir.array<8x9x?x?xf32>> {
+  %1 = arith.constant 64 : index
+  %2 = fir.alloca !fir.array<8x9x?x?xf32>, %0, %1
+  return %2 : !fir.ref<!fir.array<8x9x?x?xf32>>
+}
+
+// CHECK-LABEL: llvm.func @alloca_const_interior_array
+// CHECK-SAME: ([[OP1:%.*]]: i64) -> !llvm.ptr<array<9 x array<8 x f32>
+// CHECK: [[OP2:%.*]] = llvm.mlir.constant(64 : index) : i64
+// CHECK: [[ONE:%.*]] = llvm.mlir.constant(1 : i64) : i64
+// CHECK: [[MUL1:%.*]] = llvm.mul [[ONE]], [[OP1]] : i64
+// CHECK: [[TOTAL:%.*]] = llvm.mul [[MUL1]], [[OP2]] : i64
+// CHECK: [[A:%.*]] = llvm.alloca [[TOTAL]] x !llvm.array<9 x array<8 x f32>
+// CHECK: llvm.return [[A]] : !llvm.ptr<array<9 x array<8 x f32>
+
+// -----
+
 // Test alloca with an array with holes.
 // Constant factor of 60 (4*3*5) must be included.
 
 func @alloca_array_with_holes(%0 : index, %1 : index) -> !fir.ref<!fir.array<4x?x3x?x5xi32>> {
   %a = fir.alloca !fir.array<4x?x3x?x5xi32>, %0, %1
   return %a : !fir.ref<!fir.array<4x?x3x?x5xi32>>
 }
 

 // CHECK:         %[[C4:.*]] = llvm.mlir.constant(4 : index) : i64
 // CHECK:         %[[C1:.*]] = llvm.mlir.constant(1 : index) : i64
 // CHECK:         %[[C2:.*]] = llvm.mlir.constant(2 : index) : i64
 // CHECK:         %[[N1_TMP:.*]] = llvm.sub %[[N]], %[[SH1]]  : i64
 // CHECK:         %[[N1:.*]] = llvm.add %[[N1_TMP]], %[[C1]]  : i64
 // CHECK:         %[[N2_TMP:.*]] = llvm.sub %[[N]], %[[SH2]]  : i64
 // CHECK:         %[[N2:.*]] = llvm.add %[[N2_TMP]], %[[C1]]  : i64
 // CHECK:         %[[C1_0:.*]] = llvm.mlir.constant(1 : i64) : i64
-// CHECK:         %[[C1_1:.*]] = llvm.mlir.constant(1 : i64) : i64
-// CHECK:         %[[ARR_SIZE_TMP0:.*]] = llvm.mul %[[C1_0]], %[[C1_1]]  : i64
-// CHECK:         %[[ARR_SIZE_TMP1:.*]] = llvm.mul %[[ARR_SIZE_TMP0]], %[[N1]]  : i64
+// CHECK:         %[[ARR_SIZE_TMP1:.*]] = llvm.mul %[[C1_0]], %[[N1]]  : i64
 // CHECK:         %[[ARR_SIZE:.*]] = llvm.mul %[[ARR_SIZE_TMP1]], %[[N2]]  : i64
 // CHECK:         %[[ARR:.*]] = llvm.alloca %[[ARR_SIZE]] x f64 {bindc_name = "arr", in_type = !fir.array<?x?xf64>, operand_segment_sizes = dense<[0, 2]> : vector<2xi32>, uniq_name = "_QFsbEarr"} : (i64) -> !llvm.ptr<f64>
 // CHECK:         %[[BOX0:.*]] = llvm.mlir.undef : !llvm.struct<(ptr<f64>, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<2 x array<3 x i64>>)>
 // CHECK:         %[[ELEM_LEN:.*]] = llvm.mlir.constant(8 : i32) : i32
 // CHECK:         %[[TYPE_CODE:.*]] = llvm.mlir.constant(26 : i32) : i32
 // CHECK:         %[[ELEM_LEN_I64:.*]] = llvm.sext %[[ELEM_LEN]] : i32 to i64
 // CHECK:         %[[BOX1:.*]] = llvm.insertvalue %[[ELEM_LEN_I64]], %[[BOX0]][1 : i32] : !llvm.struct<(ptr<f64>, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<2 x array<3 x i64>>)>
 // CHECK:         %[[VERSION:.*]] = llvm.mlir.constant(20180515 : i32) : i32