[fir] Add fir ragged array builder

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.

Reviewed By: kiranchandramohan

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

Co-authored-by: Eric Schweitz <eschweitz@nvidia.com>

flang/include/flang/Optimizer/Builder/FIRBuilder.h

@@ -405,6 +405,11 @@ mlir::Value locationToFilename(fir::FirOpBuilder &, mlir::Location);
 /// Generate a constant of the given type with the location line number
 mlir::Value locationToLineNo(fir::FirOpBuilder &, mlir::Location, mlir::Type);
 
+/// Builds and returns the type of a ragged array header used to cache mask
+/// evaluations. RaggedArrayHeader is defined in
+/// flang/include/flang/Runtime/ragged.h.
+mlir::TupleType getRaggedArrayHeaderType(fir::FirOpBuilder &builder);
+
 } // namespace fir::factory
 
 #endif // FORTRAN_OPTIMIZER_BUILDER_FIRBUILDER_H

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

@@ -0,0 +1,39 @@
+//===-- 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/include/flang/Runtime/ragged.h

@@ -16,11 +16,11 @@ namespace Fortran::runtime {
 
 // A ragged array header block.
 // The header block is used to create the "array of arrays" ragged data
-// structure. It contains a pair in `flags` to indicate if the header points to
+// structure. It contains a boolean value to indicate if the header points to
 // an array of headers (isIndirection) or data elements and the rank of the
-// pointed-to array. The rank is the length of the extents vector accessed
-// through `extentPointer`. The `bufferPointer` is overloaded and is null,
-// points to an array of headers (isIndirection), or data.
+// pointed-to array in an integer value. The rank is the length of the extents
+// vector accessed through `extentPointer`. The `bufferPointer` is overloaded
+// and is null, points to an array of headers (isIndirection), or data.
 // By default, a header is set to zero, which is its unused state.
 // The layout of a ragged buffer header is mirrored in the compiler.
 struct RaggedArrayHeader {

flang/lib/Optimizer/Builder/CMakeLists.txt

@@ -11,6 +11,7 @@ add_flang_library(FIRBuilder
   Runtime/Character.cpp
   Runtime/Derived.cpp
   Runtime/Numeric.cpp
+  Runtime/Ragged.cpp
   Runtime/Reduction.cpp
   Runtime/Transformational.cpp
 

flang/lib/Optimizer/Builder/FIRBuilder.cpp

@@ -594,3 +594,15 @@ fir::factory::createExtents(fir::FirOpBuilder &builder, mlir::Location loc,
             : builder.createIntegerConstant(loc, idxTy, ext));
   return extents;
 }
+
+mlir::TupleType
+fir::factory::getRaggedArrayHeaderType(fir::FirOpBuilder &builder) {
+  mlir::IntegerType i1Ty = builder.getIntegerType(1);
+  mlir::IntegerType i8Ty = builder.getIntegerType(8);
+  mlir::IntegerType i64Ty = builder.getIntegerType(64);
+  auto arrTy = fir::SequenceType::get(builder.getIntegerType(8), 1);
+  auto buffTy = fir::HeapType::get(arrTy);
+  auto extTy = fir::SequenceType::get(i64Ty, 1);
+  auto shTy = fir::HeapType::get(extTy);
+  return mlir::TupleType::get(builder.getContext(), {i1Ty, i8Ty, buffTy, shTy});
+}

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

@@ -0,0 +1,67 @@
+//===-- 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);
+  // Position of the bufferPointer in the header struct.
+  auto two = builder.createIntegerConstant(loc, i32Ty, 2);
+  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, two);
+  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

@@ -0,0 +1,36 @@
+//===- 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::TupleType headerTy =
+      fir::factory::getRaggedArrayHeaderType(*firBuilder);
+  mlir::Value header = firBuilder->create<fir::UndefOp>(loc, headerTy);
+  mlir::Value eleSize = firBuilder->createIntegerConstant(loc, i32Ty, 1);
+  mlir::Value extent = firBuilder->createIntegerConstant(loc, i32Ty, 1);
+  // Use a dummy header just to test the correctness of the generated call.
+  fir::runtime::genRaggedArrayAllocate(
+      loc, *firBuilder, header, false, eleSize, {extent});
+  checkCallOpFromResultBox(
+      eleSize, "_FortranARaggedArrayAllocate", 5, /*addLocArgs=*/false);
+}
+
+TEST_F(RuntimeCallTest, genRaggedArrayDeallocateTest) {
+  auto loc = firBuilder->getUnknownLoc();
+  mlir::TupleType headerTy =
+      fir::factory::getRaggedArrayHeaderType(*firBuilder);
+  // Use a dummy header just to test the correctness of the generated call.
+  mlir::Value header = firBuilder->create<fir::UndefOp>(loc, headerTy);
+  fir::runtime::genRaggedArrayDeallocate(loc, *firBuilder, header);
+  checkCallOpFromResultBox(
+      header, "_FortranARaggedArrayDeallocate", 1, /*addLocArgs=*/false);
+}

flang/unittests/Optimizer/CMakeLists.txt

@@ -17,6 +17,7 @@ add_flang_unittest(FlangOptimizerTests
   Builder/Runtime/CharacterTest.cpp
   Builder/Runtime/DerivedTest.cpp
   Builder/Runtime/NumericTest.cpp
+  Builder/Runtime/RaggedTest.cpp
   Builder/Runtime/ReductionTest.cpp
   Builder/Runtime/TransformationalTest.cpp
   FIRContextTest.cpp