Merge pull request #1109 from schweitzpgi/ch-speicher

Full implementation of lazy, ragged buffers for caching intermediate

Author: schweitzpgi

flang/include/flang/Lower/ConvertExpr.h

@@ -169,28 +169,19 @@ createSomeArrayTempValue(AbstractConverter &converter,
                          const evaluate::Expr<evaluate::SomeType> &expr,
                          SymMap &symMap, StatementContext &stmtCtx);
 
-// Lambda to reload the dynamically allocated pointers to a lazy buffer and its
-// extents. This is used to introduce these ssa-values in a place that will
-// dominate any/all subsequent uses after the loop that created the lazy buffer.
-using LoadLazyBufferLambda =
-    std::function<std::pair<fir::ExtendedValue, mlir::Value>(
-        fir::FirOpBuilder &)>;
-
-// Creating a lazy array temporary returns a pair of values. The first is an
-// extended value which is a pointer to the buffer, of array type, with the
-// appropriate dynamic extents. The second argument is a continuation to reload
-// the buffer at some future point in the code gen.
-using CreateLazyArrayResult =
-    std::pair<fir::ExtendedValue, LoadLazyBufferLambda>;
-
-/// Like createSomeArrayTempValue, but the temporary buffer is allocated lazily
-/// (inside the loops instead of before the loops). This can be useful if a
-/// loop's bounds are functions of other loop indices, for example.
-CreateLazyArrayResult
-createLazyArrayTempValue(AbstractConverter &converter,
-                         const evaluate::Expr<evaluate::SomeType> &expr,
-                         mlir::Value var, mlir::Value shapeBuffer,
-                         SymMap &symMap, StatementContext &stmtCtx);
+/// Somewhat similar to createSomeArrayTempValue, but the temporary buffer is
+/// allocated lazily (inside the loops instead of before the loops) to
+/// accomodate buffers with shapes that cannot be precomputed. In fact, the
+/// buffer need not even be hyperrectangular. The buffer may be created as an
+/// instance of a ragged array, which may be useful if an array's extents are
+/// functions of other loop indices. The ragged array structure is built with \p
+/// raggedHeader being the root header variable. The header is a tuple of
+/// `{rank, data-is-headers, [data]*, [extents]*}`, which is built recursively.
+/// The base header, \p raggedHeader, must be initialized to zeros.
+void createLazyArrayTempValue(AbstractConverter &converter,
+                              const evaluate::Expr<evaluate::SomeType> &expr,
+                              mlir::Value raggedHeader, SymMap &symMap,
+                              StatementContext &stmtCtx);
 
 /// Lower an array expression to a value of type box. The expression must be a
 /// variable.

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

@@ -465,6 +465,8 @@ void genRecordAssignment(fir::FirOpBuilder &builder, mlir::Location loc,
                          const fir::ExtendedValue &lhs,
                          const fir::ExtendedValue &rhs);
 
+mlir::TupleType getRaggedArrayHeaderType(fir::FirOpBuilder &builder);
+
 } // namespace fir::factory
 
 #endif // FORTRAN_OPTIMIZER_BUILDER_FIRBUILDER_H

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

@@ -105,10 +105,16 @@ constexpr TypeBuilderFunc getModel<signed char>() {
   };
 }
 template <>
+constexpr TypeBuilderFunc getModel<void *>() {
+  return [](mlir::MLIRContext *context) -> mlir::Type {
+    return fir::LLVMPointerType::get(mlir::IntegerType::get(context, 8));
+  };
+}
+template <>
 constexpr TypeBuilderFunc getModel<void **>() {
   return [](mlir::MLIRContext *context) -> mlir::Type {
     return fir::ReferenceType::get(
-        fir::PointerType::get(mlir::IntegerType::get(context, 8)));
+        fir::LLVMPointerType::get(mlir::IntegerType::get(context, 8)));
   };
 }
 template <>
@@ -125,6 +131,10 @@ constexpr TypeBuilderFunc getModel<long &>() {
   };
 }
 template <>
+constexpr TypeBuilderFunc getModel<long *>() {
+  return getModel<long &>();
+}
+template <>
 constexpr TypeBuilderFunc getModel<long long>() {
   return [](mlir::MLIRContext *context) -> mlir::Type {
     return mlir::IntegerType::get(context, 8 * sizeof(std::size_t));
@@ -138,6 +148,10 @@ constexpr TypeBuilderFunc getModel<long long &>() {
   };
 }
 template <>
+constexpr TypeBuilderFunc getModel<long long *>() {
+  return getModel<long long &>();
+}
+template <>
 constexpr TypeBuilderFunc getModel<unsigned long>() {
   return [](mlir::MLIRContext *context) -> mlir::Type {
     return mlir::IntegerType::get(context, 8 * sizeof(unsigned long));

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

@@ -0,0 +1,24 @@
+//===-- Runtime/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_RUNTIME_RAGGED_H_
+#define FORTRAN_RUNTIME_RAGGED_H_
+
+#include "flang/Runtime/entry-names.h"
+#include <cstdint>
+
+namespace Fortran::runtime {
+extern "C" {
+// Helper for allocation of ragged array buffer blocks.
+void *RTNAME(RaggedArrayAllocate)(
+    void *, bool, std::int64_t, std::int64_t, std::int64_t *);
+// Helper for deallocation of ragged array buffers.
+void RTNAME(RaggedArrayDeallocate)(void *);
+} // extern "C"
+} // namespace Fortran::runtime
+#endif // FORTRAN_RUNTIME_RAGGED_H_

flang/lib/Lower/Bridge.cpp

@@ -32,6 +32,7 @@
 #include "flang/Optimizer/Builder/Character.h"
 #include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Builder/Runtime/Character.h"
+#include "flang/Optimizer/Builder/Runtime/Ragged.h"
 #include "flang/Optimizer/Dialect/FIRAttr.h"
 #include "flang/Optimizer/Dialect/FIRDialect.h"
 #include "flang/Optimizer/Dialect/FIROps.h"
@@ -1263,8 +1264,9 @@ class FirConverter : public Fortran::lower::AbstractConverter {
         forceControlVariableBinding(ctrlVar, lp.getInductionVar());
         loops.push_back(lp);
       }
-      if (outermost)
-         explicitIterSpace.setOuterLoop(loops[0]);
+      if (outermost) 
+        explicitIterSpace.setOuterLoop(loops[0]);
+      explicitIterSpace.appendLoops(loops);
       if (const auto &mask =
               std::get<std::optional<Fortran::parser::ScalarLogicalExpr>>(
                   header.t);
@@ -2749,23 +2751,36 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     // with sequences of i1. That is, an array of i1 will be truncated in size
     // and be too small. For example, a buffer of type fir.array<7xi1> will have
     // 0 size.
-    auto ty = fir::HeapType::get(builder->getIntegerType(8));
+    auto i64Ty = builder->getIntegerType(64);
+    auto ty = fir::factory::getRaggedArrayHeaderType(*builder);
+    auto buffTy = ty.getType(1);
+    auto shTy = ty.getType(2);
     auto loc = toLocation();
-    auto var = builder->createTemporary(loc, ty);
-    auto nil = builder->createNullConstant(loc, ty);
-    builder->create<fir::StoreOp>(loc, nil, var);
-    auto shTy = fir::HeapType::get(builder->getIndexType());
-    auto shape = builder->createTemporary(loc, shTy);
-    auto nilSh = builder->createNullConstant(loc, shTy);
-    builder->create<fir::StoreOp>(loc, nilSh, shape);
-    implicitIterSpace.addMaskVariable(exp, var, shape);
+    auto hdr = builder->createTemporary(loc, ty);
+    // FIXME: Is there a way to create a `zeroinitializer` in LLVM-IR dialect?
+    // For now, explicitly set lazy ragged header to all zeros.
+    // auto nilTup = builder->createNullConstant(loc, ty);
+    // builder->create<fir::StoreOp>(loc, nilTup, hdr);
+    auto i32Ty = builder->getIntegerType(32);
+    auto zero = builder->createIntegerConstant(loc, i32Ty, 0);
+    auto zero64 = builder->createIntegerConstant(loc, i64Ty, 0);
+    mlir::Value flags = builder->create<fir::CoordinateOp>(
+        loc, builder->getRefType(i64Ty), hdr, zero);
+    builder->create<fir::StoreOp>(loc, zero64, flags);
+    auto one = builder->createIntegerConstant(loc, i32Ty, 1);
+    auto nullPtr1 = builder->createNullConstant(loc, buffTy);
+    mlir::Value var = builder->create<fir::CoordinateOp>(
+        loc, builder->getRefType(buffTy), hdr, one);
+    builder->create<fir::StoreOp>(loc, nullPtr1, var);
+    auto two = builder->createIntegerConstant(loc, i32Ty, 2);
+    auto nullPtr2 = builder->createNullConstant(loc, shTy);
+    mlir::Value shape = builder->create<fir::CoordinateOp>(
+        loc, builder->getRefType(shTy), hdr, two);
+    builder->create<fir::StoreOp>(loc, nullPtr2, shape);
+    implicitIterSpace.addMaskVariable(exp, var, shape, hdr);
     explicitIterSpace.outermostContext().attachCleanup(
-        [builder = this->builder, loc, var]() {
-          auto load = builder->create<fir::LoadOp>(loc, var);
-          auto cmp = builder->genIsNotNull(loc, load);
-          builder->genIfThen(loc, cmp)
-              .genThen([&]() { builder->create<fir::FreeMemOp>(loc, load); })
-              .end();
+        [builder = this->builder, hdr, loc]() {
+          fir::runtime::genRaggedArrayDeallocate(loc, *builder, hdr);
         });
   }