[flang] Fix AllocaOp/AllocMemOp type conversion

[peixin]

Summary

For arrays without a constant interior or arrays of character with
dynamic length arrays, the data types are converted to a pointer to the
element type, so the scale size of the constant extents needs to be
counted. The previous AllocaOp conversion does not consider the arrays
of character with dynamic length arrays, and the previous AllocMemOp
conversion does not consider arrays without a constant interior. This
fixes them and refactors the code so that it can be shared. Also add
the test cases.

[jeanPerier]

Thanks for the bug fix

flang/lib/Optimizer/CodeGen/CodeGen.cpp

354

We should really try to share this part of the code with fir.allocmem codegen (that may or may not contain this bug, but that is doing very similar logic here: https://github.com/llvm/llvm-project/blob/44e58509be72b49c3abc7b4c5a7ffa81a754fe1c/flang/lib/Optimizer/CodeGen/CodeGen.cpp#L1006). Can you try to share it (a lot more can probably be shared between allocmem and alloca codegen) ?

[peixin]

Thanks @jeanPerier for the notice. The AllocMemOp does have one similar bug. Also fix it and share the code with AllocaOp.

[jeanPerier]

Thanks a lot @peixin, it is great to have this code shared and to have more regression tests here !

[peixin]

Thanks @jeanPerier . FIR upstream work is not done. So, do I need to cherry-pick and sync this, D124345, and D125022 to fir-dev after landing them considering it may affect the upstream work?

flang/lib/Optimizer/CodeGen/CodeGen.cpp

matchAndRewrite(fir::AllocaOp alloc, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     mlir::ValueRange operands = adaptor.getOperands();
     auto loc = alloc.getLoc();
     mlir::Type ity = lowerTy().indexType();
     unsigned i = 0;
     mlir::Value size = genConstantIndex(loc, ity, rewriter, 1).getResult();
     mlir::Type ty = convertType(alloc.getType());
     mlir::Type resultTy = ty;
+    mlir::Type allocEleTy = alloc.getInType();
+    mlir::Type scalarType = fir::unwrapSequenceType(allocEleTy);
     if (alloc.hasLenParams()) {
       unsigned end = alloc.numLenParams();
       llvm::SmallVector<mlir::Value> lenParams;
       for (; i < end; ++i)
         lenParams.push_back(operands[i]);
-      mlir::Type scalarType = fir::unwrapSequenceType(alloc.getInType());
       if (auto chrTy = scalarType.dyn_cast<fir::CharacterType>()) {
         fir::CharacterType rawCharTy = fir::CharacterType::getUnknownLen(
             chrTy.getContext(), chrTy.getFKind());
         ty = mlir::LLVM::LLVMPointerType::get(convertType(rawCharTy));
         assert(end == 1);
         size = integerCast(loc, rewriter, ity, lenParams[0]);
       } else if (auto recTy = scalarType.dyn_cast<fir::RecordType>()) {
         mlir::LLVM::LLVMFuncOp memSizeFn =
             getDependentTypeMemSizeFn(recTy, alloc, rewriter);
         if (!memSizeFn)
           emitError(loc, "did not find allocation function");
         mlir::NamedAttribute attr = rewriter.getNamedAttr(
             "callee", mlir::SymbolRefAttr::get(memSizeFn));
         auto call = rewriter.create<mlir::LLVM::CallOp>(
             loc, ity, lenParams, llvm::ArrayRef<mlir::NamedAttribute>{attr});
         size = call.getResult(0);
         ty = mlir::LLVM::LLVMPointerType::get(
             mlir::IntegerType::get(alloc.getContext(), 8));
       } 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.
+    // We do this for char arrays and non-char arrays without a constant
+    // interior, since those are the only ones that get decayed to a pointer
+    // to the element type.
+    if ((allocEleTy.isa<fir::SequenceType>() &&
+         scalarType.isa<fir::CharacterType>()) ||
+        alloc.hasShapeOperands()) {
       if (auto seqTy = allocEleTy.dyn_cast<fir::SequenceType>()) {
-        if (!seqTy.hasConstantInterior()) {
+        if (scalarType.isa<fir::CharacterType>() ||
+            !seqTy.hasConstantInterior()) {
           fir::SequenceType::Extent constSize = 1;
           for (auto extent : seqTy.getShape())
             if (extent != fir::SequenceType::getUnknownExtent())
               constSize *= extent;
+          if (constSize != 1) {
             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/alloc.fir

func @char_array_alloca(%l: i32, %e : index) -> !fir.ref<!fir.array<?x?x!fir.char<1,?>>> {
   // CHECK: %[[lcast:.*]] = sext i32 %[[l]] to i64
   // CHECK: %[[prod:.*]] = mul i64 %[[lcast]], %[[e]]
   // CHECK: %[[size:.*]] = mul i64 %[[prod]], %[[e]]
   // CHECK: alloca i8, i64 %[[size]]
   %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: define i8* @char_array_alloca2(
+// CHECK-SAME: i32 %[[l:.*]], i64 %[[e:.*]])
+func @char_array_alloca2(%l: i32, %e : index) -> !fir.ref<!fir.array<3x?x!fir.char<1,?>>> {
+  // CHECK: %[[lcast:.*]] = sext i32 %[[l]] to i64
+  // CHECK: %[[prod:.*]] = mul i64 %[[lcast]], 3
+  // CHECK: %[[size:.*]] = mul i64 %[[prod]], %[[e]]
+  // CHECK: alloca i8, i64 %[[size]]
+  %a = fir.alloca !fir.array<3x?x!fir.char<1,?>>(%l : i32), %e
+  return %a :  !fir.ref<!fir.array<3x?x!fir.char<1,?>>>
+}
+
+// CHECK-LABEL: define i8* @char_array_alloca3(
+// CHECK-SAME: i32 %[[l:.*]])
+func @char_array_alloca3(%l: i32) -> !fir.ref<!fir.array<3x3x!fir.char<1,?>>> {
+  // CHECK: %[[lcast:.*]] = sext i32 %[[l]] to i64
+  // CHECK: %[[size:.*]] = mul i64 %[[lcast]], 9
+  // CHECK: alloca i8, i64 %[[size]]
+  %a = fir.alloca !fir.array<3x3x!fir.char<1,?>>(%l : i32)
+  return %a :  !fir.ref<!fir.array<3x3x!fir.char<1,?>>>
+}
+
 // Constant factor of 60 (4*3*5) must be included.
 // CHECK-LABEL: define i32* @array_with_holes(
 // CHECK-SAME: i64 %[[a:.*]], i64 %[[b:.*]])
 func @array_with_holes(%0 : index, %1 : index) -> !fir.ref<!fir.array<4x?x3x?x5xi32>> {
   // CHECK: %[[prod1:.*]] = mul i64 60, %[[a]]
   // CHECK: %[[prod2:.*]] = mul i64 %[[prod1]], %[[b]]
   // CHECK: alloca i32, i64 %[[prod2]]
   %a = fir.alloca !fir.array<4x?x3x?x5xi32>, %0, %1