[fir] Add fir ragged array builder

[clementval]

Summary

This patch introduces a bunch of builder functions
to create function calls to runtime ragged arrays functions.

This patch is part of the upstreaming effort from fir-dev branch.

Co-authored-by: Eric Schweitz <[e-mail]>

[clementval]

Rebase on top of D114557

[kiranchandramohan]

I have a couple of questions.

flang/lib/Optimizer/Builder/Runtime/Ragged.cpp

35

I guess this is two now, right?

flang/unittests/Optimizer/Builder/Runtime/RaggedTest.cpp

19

Nit: I guess this is not a true representation of a header.
But I guess it is OK since we are not testing the structure of the header.

Also, is the tupleTy/pair due to the previous representation of the raggedArrayHeader?

32

Nit: This seems to be different from the header creation above.

I guess a common getHeader function to get the header and a comment indicating that this is not a true header but just used for testing the call to the Ragged runtime functions might be helpful.

[clementval]

Address review comments

flang/lib/Optimizer/Builder/Runtime/Ragged.cpp

35

Good catch. Updated and I added a comment for this magic number.

flang/unittests/Optimizer/Builder/Runtime/RaggedTest.cpp

19

There is actually a helper function in the FIRBuilder for that so I just bring it in with this patch.

[kiranchandramohan]

LGTM.

flang/include/flang/Optimizer/Builder/Runtime/Ragged.h

+//===-- Ragged.h ------------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef FORTRAN_OPTIMIZER_BUILDER_RUNTIME_RAGGED_H
+#define FORTRAN_OPTIMIZER_BUILDER_RUNTIME_RAGGED_H
+
+namespace mlir {
+class Location;
+class Value;
+class ValueRange;
+} // namespace mlir
+
+namespace fir {
+class FirOpBuilder;
+} // namespace fir
+
+namespace fir::runtime {
+
+/// Generate code to instantiate a section of a ragged array. Calls the runtime
+/// to initialize the data buffer. \p header must be a ragged buffer header (on
+/// the heap) and will be initialized, if and only if the rank of \p extents is
+/// at least 1 and all values in the vector of extents are positive. \p extents
+/// must be a vector of Value of type `i64`. \p eleSize is in bytes, not bits.
+void genRaggedArrayAllocate(mlir::Location loc, fir::FirOpBuilder &builder,
+                            mlir::Value header, bool asHeaders,
+                            mlir::Value eleSize, mlir::ValueRange extents);
+
+/// Generate a call to the runtime routine to deallocate a ragged array data
+/// structure on the heap.
+void genRaggedArrayDeallocate(mlir::Location loc, fir::FirOpBuilder &builder,
+                              mlir::Value header);
+
+} // namespace fir::runtime
+#endif // FORTRAN_OPTIMIZER_BUILDER_RUNTIME_RAGGED_H

flang/lib/Optimizer/Builder/CMakeLists.txt

 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 
 add_flang_library(FIRBuilder
   BoxValue.cpp
   Character.cpp
   Complex.cpp
   DoLoopHelper.cpp
   FIRBuilder.cpp
   MutableBox.cpp
+  Runtime/Ragged.cpp
 
   DEPENDS
   FIRDialect
   FIRSupport
   ${dialect_libs}
 
   LINK_LIBS
   FIRDialect
   FIRSupport
   ${dialect_libs}
 )

flang/lib/Optimizer/Builder/Runtime/Ragged.cpp

+//===-- Ragged.cpp --------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Builder/Runtime/Ragged.h"
+#include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "flang/Optimizer/Builder/Runtime/RTBuilder.h"
+#include "flang/Runtime/ragged.h"
+
+using namespace Fortran::runtime;
+
+void fir::runtime::genRaggedArrayAllocate(mlir::Location loc,
+                                          fir::FirOpBuilder &builder,
+                                          mlir::Value header, bool asHeaders,
+                                          mlir::Value eleSize,
+                                          mlir::ValueRange extents) {
+  auto i32Ty = builder.getIntegerType(32);
+  auto rank = extents.size();
+  auto i64Ty = builder.getIntegerType(64);
+  auto func =
+      fir::runtime::getRuntimeFunc<mkRTKey(RaggedArrayAllocate)>(loc, builder);
+  auto fTy = func.getType();
+  auto i1Ty = builder.getIntegerType(1);
+  fir::SequenceType::Shape shape = {
+      static_cast<fir::SequenceType::Extent>(rank)};
+  auto extentTy = fir::SequenceType::get(shape, i64Ty);
+  auto refTy = fir::ReferenceType::get(i64Ty);
+  auto one = builder.createIntegerConstant(loc, i32Ty, 1);
+  auto eleTy = fir::unwrapSequenceType(fir::unwrapRefType(header.getType()));
+  auto ptrTy = builder.getRefType(eleTy.cast<mlir::TupleType>().getType(1));
+  auto ptr = builder.create<fir::CoordinateOp>(loc, ptrTy, header, one);
+  auto heap = builder.create<fir::LoadOp>(loc, ptr);
+  auto cmp = builder.genIsNull(loc, heap);
+  builder.genIfThen(loc, cmp)
+      .genThen([&]() {
+        auto asHeadersVal = builder.createIntegerConstant(loc, i1Ty, asHeaders);
+        auto rankVal = builder.createIntegerConstant(loc, i64Ty, rank);
+        auto buff = builder.create<fir::AllocMemOp>(loc, extentTy);
+        // Convert all the extents to i64 and pack them in a buffer on the heap.
+        for (auto i : llvm::enumerate(extents)) {
+          auto offset = builder.createIntegerConstant(loc, i32Ty, i.index());
+          auto addr =
+              builder.create<fir::CoordinateOp>(loc, refTy, buff, offset);
+          auto castVal = builder.createConvert(loc, i64Ty, i.value());
+          builder.create<fir::StoreOp>(loc, castVal, addr);
+        }
+        auto args = fir::runtime::createArguments(
+            builder, loc, fTy, header, asHeadersVal, rankVal, eleSize, buff);
+        builder.create<fir::CallOp>(loc, func, args);
+      })
+      .end();
+}
+
+void fir::runtime::genRaggedArrayDeallocate(mlir::Location loc,
+                                            fir::FirOpBuilder &builder,
+                                            mlir::Value header) {
+  auto func = fir::runtime::getRuntimeFunc<mkRTKey(RaggedArrayDeallocate)>(
+      loc, builder);
+  auto fTy = func.getType();
+  auto args = fir::runtime::createArguments(builder, loc, fTy, header);
+  builder.create<fir::CallOp>(loc, func, args);
+}

flang/unittests/Optimizer/Builder/Runtime/RaggedTest.cpp

+//===- RaggedTest.cpp -- Ragged array runtime function builder unit tests -===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Builder/Runtime/Ragged.h"
+#include "RuntimeCallTestBase.h"
+#include "gtest/gtest.h"
+
+TEST_F(RuntimeCallTest, genRaggedArrayAllocateTest) {
+  auto loc = firBuilder->getUnknownLoc();
+  mlir::Type tupleTy =
+      mlir::TupleType::get(firBuilder->getContext(), {i32Ty, i32Ty});
+  mlir::Type seqTy =
+      fir::SequenceType::get(fir::SequenceType::Shape(1, 10), tupleTy);
+  mlir::Value header = firBuilder->create<fir::UndefOp>(loc, seqTy);
+  mlir::Value eleSize = firBuilder->createIntegerConstant(loc, i32Ty, 1);
+  mlir::Value extent = firBuilder->createIntegerConstant(loc, i32Ty, 1);
+  fir::runtime::genRaggedArrayAllocate(
+      loc, *firBuilder, header, false, eleSize, {extent});
+  checkCallOpFromResultBox(
+      eleSize, "_FortranARaggedArrayAllocate", 5, /*addLocArgs=*/false);
+}
+
+TEST_F(RuntimeCallTest, genRaggedArrayDeallocateTest) {
+  auto loc = firBuilder->getUnknownLoc();
+  mlir::Type seqTy =
+      fir::SequenceType::get(fir::SequenceType::Shape(1, 10), i32Ty);
+  mlir::Value header = firBuilder->create<fir::UndefOp>(loc, seqTy);
+  fir::runtime::genRaggedArrayDeallocate(loc, *firBuilder, header);
+  checkCallOpFromResultBox(
+      header, "_FortranARaggedArrayDeallocate", 1, /*addLocArgs=*/false);
+}

flang/unittests/Optimizer/CMakeLists.txt

 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 
 set(LIBS
   FIRBuilder
   FIRCodeGen
   FIRDialect
   FIRSupport
   ${dialect_libs}
 )
 
 add_flang_unittest(FlangOptimizerTests
   Builder/CharacterTest.cpp
   Builder/ComplexTest.cpp
   Builder/DoLoopHelperTest.cpp
   Builder/FIRBuilderTest.cpp
+  Builder/Runtime/RaggedTest.cpp
   FIRContextTest.cpp
   InternalNamesTest.cpp
   KindMappingTest.cpp
   RTBuilder.cpp
 )
 target_link_libraries(FlangOptimizerTests
   PRIVATE
   ${LIBS})