[flang] Upstream changes to allocmem codegen

[schweitz]

Summary

Fix a correctness issue with arrays of characters. Include tests, etc.

[clementval]

LGTM

[schweitz]

flang doesn't fully support windows yet.

[schweitz]

Fix comment

[PeteSteinfeld]

All builds and tests correctly and looks good.

flang/include/flang/Optimizer/Dialect/FIRType.h

 #ifndef NDEBUG
 // !fir.ptr<X> and !fir.heap<X> where X is !fir.ptr, !fir.heap, or !fir.ref
 // is undefined and disallowed.
 inline bool singleIndirectionLevel(mlir::Type ty) {
   return !fir::isa_ref_type(ty);
 }
 #endif
 
+/// Return true iff `ty` is a RecordType with type parameters.
+inline bool isRecordWithTypeParameters(mlir::Type ty) {
+  if (auto recTy = ty.dyn_cast_or_null<fir::RecordType>())
+    return recTy.getNumLenParams() != 0;
+  return false;
+}
+
 /// Apply the components specified by `path` to `rootTy` to determine the type
 /// of the resulting component element. `rootTy` should be an aggregate type.
 /// Returns null on error.
 mlir::Type applyPathToType(mlir::Type rootTy, mlir::ValueRange path);
 
 } // namespace fir
 
 #endif // FORTRAN_OPTIMIZER_DIALECT_FIRTYPE_H

flang/lib/Optimizer/CodeGen/CodeGen.cpp

 namespace {
 /// Lower a `fir.allocmem` instruction into `llvm.call @malloc`
 struct AllocMemOpConversion : public FIROpConversion<fir::AllocMemOp> {
   using FIROpConversion::FIROpConversion;
 
   mlir::LogicalResult
   matchAndRewrite(fir::AllocMemOp heap, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
-    mlir::Type ty = convertType(heap.getType());
+    auto heapTy = heap.getType();
+    auto ty = convertType(heapTy);
     mlir::LLVM::LLVMFuncOp mallocFunc = getMalloc(heap, rewriter);
     mlir::Location loc = heap.getLoc();
     auto ity = lowerTy().indexType();
-    if (auto recTy = fir::unwrapSequenceType(heap.getAllocatedType())
-                         .dyn_cast<fir::RecordType>())
-      if (recTy.getNumLenParams() != 0) {
-        TODO(loc,
-             "fir.allocmem codegen of derived type with length parameters");
-        return failure();
-      }
+    auto dataTy = fir::unwrapRefType(heapTy);
+    if (fir::isRecordWithTypeParameters(fir::unwrapSequenceType(dataTy)))
+      TODO(loc, "fir.allocmem codegen of derived type with length parameters");
     mlir::Value size = genTypeSizeInBytes(loc, ity, rewriter, ty);
+    // !fir.array<NxMx!fir.char<K,?>> sets `size` to the width of !fir.char<K>.
+    // So multiply the constant dimensions here.
+    if (fir::hasDynamicSize(dataTy))
+      if (auto seqTy = dataTy.dyn_cast<fir::SequenceType>())
+        if (fir::characterWithDynamicLen(seqTy.getEleTy())) {
+          fir::SequenceType::Extent arrSize = 1;
+          for (auto d : seqTy.getShape())
+            if (d != fir::SequenceType::getUnknownExtent())
+              arrSize *= d;
+          size = rewriter.create<mlir::LLVM::MulOp>(
+              loc, ity, size, genConstantIndex(loc, ity, rewriter, arrSize));
+        }
     for (mlir::Value opnd : adaptor.getOperands())
       size = rewriter.create<mlir::LLVM::MulOp>(
           loc, ity, size, integerCast(loc, rewriter, ity, opnd));
     heap->setAttr("callee", mlir::SymbolRefAttr::get(mallocFunc));
     auto malloc = rewriter.create<mlir::LLVM::CallOp>(
         loc, ::getVoidPtrType(heap.getContext()), size, heap->getAttrs());
     rewriter.replaceOpWithNewOp<mlir::LLVM::BitcastOp>(heap, ty,
                                                        malloc.getResult(0));

flang/test/Fir/alloc.fir

+// RUN: tco %s | FileCheck %s
+
+// UNSUPPORTED: system-windows
+
+// CHECK-LABEL: define i32* @f1()
+func @f1() -> !fir.ref<i32> {
+  // CHECK: alloca i32, i64 1
+  %1 = fir.alloca i32
+  return %1 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: define i32* @f2()
+func @f2() -> !fir.ref<i32> {
+  %0 = arith.constant 100 : index
+  // CHECK: alloca i32, i64 100
+  %1 = fir.alloca i32, %0
+  return %1 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: define i32* @f3()
+func @f3() -> !fir.heap<i32> {
+  // CHECK: call i8* @malloc(i64 4)
+  %1 = fir.allocmem i32
+  return %1 : !fir.heap<i32>
+}
+
+// CHECK-LABEL: define i32* @f4()
+func @f4() -> !fir.heap<i32> {
+  %0 = arith.constant 100 : index
+  // CHECK: call i8* @malloc(i64 400)
+  %1 = fir.allocmem i32, %0
+  return %1 : !fir.heap<i32>
+}
+
+// CHECK-LABEL: define i32** @f5()
+func @f5() -> !fir.ref<!fir.ptr<!fir.array<?xi32>>> {
+  // CHECK: alloca i32*, i64 1
+  %1 = fir.alloca !fir.ptr<!fir.array<?xi32>>
+  return %1 : !fir.ref<!fir.ptr<!fir.array<?xi32>>>
+}
+
+// CHECK-LABEL: define i8* @char_array_alloca(
+// CHECK-SAME: i32 %[[l:.*]], i64 %[[e:.*]])
+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,?>>>
+}
+
+// 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
+  return %a : !fir.ref<!fir.array<4x?x3x?x5xi32>>
+}
+
+// CHECK-LABEL: define void @allocmem_array_of_dynchar(
+// CHECK-SAME: i64 %[[arg:.*]])
+// CHECK: %[[mul:.*]] = mul i64 9, %[[arg]]
+// CHECK: %[[malloc:.*]] = call i8* @malloc(i64 %[[mul]])
+// CHECK: ret void
+func @allocmem_array_of_dynchar(%arg0: index) {
+  %1 = fir.allocmem !fir.array<3x3x!fir.char<1,?>>(%arg0 : index)
+  return
+}
+
+// CHECK-LABEL: define void @allocmem_dynarray_of_dynchar(
+// CHECK-SAME: i64 %[[len:.*]], i64 %[[extent:.*]])
+// CHECK: %[[a:.*]] = mul i64 24, %[[len]]
+// CHECK: %[[b:.*]] = mul i64 %[[a]], %[[extent]]
+// CHECK: %[[malloc:.*]] = call i8* @malloc(i64 %[[b]])
+// CHECK: ret void
+func @allocmem_dynarray_of_dynchar(%arg0: index, %arg1: index) {
+  %1 = fir.allocmem !fir.array<3x?x4x!fir.char<2,?>>(%arg0 : index), %arg1
+  return
+}

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

func @test_string_with_shape(%len: index, %nelems: index) {
   %1 = fir.allocmem !fir.array<?x!fir.char<1,?>>(%len : index), %nelems
   fir.freemem %1 : !fir.heap<!fir.array<?x!fir.char<1,?>>>
   return
 }
 
 // CHECK-LABEL: llvm.func @test_string_with_shape
 // CHECK-SAME: %[[LEN:.*]]: i64, %[[NELEMS:.*]]: i64)
 // CHECK:   %[[ONE:.*]] = llvm.mlir.constant(1 : i64) : i64
-// CHECK:   %[[LEN_SIZE:.*]] = llvm.mul %[[ONE]], %[[LEN]]  : i64
+// CHECK:   %[[ONE2:.*]] = llvm.mlir.constant(1 : i64) : i64
+// CHECK:   %[[MUL1:.*]] = llvm.mul %[[ONE]], %[[ONE2]]  : i64
+// CHECK:   %[[LEN_SIZE:.*]] = llvm.mul %[[MUL1]], %[[LEN]]  : i64
 // CHECK:   %[[TOTAL_SIZE:.*]] = llvm.mul %[[LEN_SIZE]], %[[NELEMS]]  : i64
 // CHECK:   %[[MEM:.*]] = llvm.call @malloc(%[[TOTAL_SIZE]])
 // CHECK:   %[[B1:.*]] = llvm.bitcast %[[MEM]] : !llvm.ptr<i8> to !llvm.ptr<i8>
 // CHECK:   %[[B2:.*]] = llvm.bitcast %[[B1]] : !llvm.ptr<i8> to !llvm.ptr<i8>
 // CHECK:   llvm.call @free(%[[B2]]) : (!llvm.ptr<i8>) -> ()
 // CHECK:   llvm.return
 // CHECK: }