[flang] Add allocator_idx attribute on fir.embox and fircg.ext_embox (#101212)

clementval · web-flow · commit 0def9a923dad · 2024-08-01T12:49:17.000-07:00
#100690 introduces allocator registry with the ability to store allocator index in the descriptor. This patch adds an attribute to fir.embox and fircg.ext_embox to be able to set the allocator index while populating the descriptor fields.
diff --git a/flang/include/flang/Optimizer/CodeGen/CGOps.td b/flang/include/flang/Optimizer/CodeGen/CGOps.td
@@ -48,6 +48,7 @@ def fircg_XEmboxOp : fircg_Op<"ext_embox", [AttrSizedOperandSegments]> {
        - substring: A substring operator (offset, length) for CHARACTER.
        - LEN type parameters: A vector of runtime LEN type parameters that
          describe an correspond to the elemental derived type.
+       - allocator_idx: specify special allocator to use.
 
     The memref and shape arguments are mandatory. The rest are optional.
   }];
@@ -60,7 +61,8 @@ def fircg_XEmboxOp : fircg_Op<"ext_embox", [AttrSizedOperandSegments]> {
     Variadic<AnyCoordinateType>:$subcomponent,
     Variadic<AnyIntegerType>:$substr,
     Variadic<AnyIntegerType>:$lenParams,
-    Optional<fir_ClassType>:$sourceBox
+    Optional<fir_ClassType>:$sourceBox,
+    OptionalAttr<I32Attr>:$allocator_idx
   );
   let results = (outs BoxOrClassType);
 
diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -783,6 +783,7 @@ def fir_EmboxOp : fir_Op<"embox", [NoMemoryEffect, AttrSizedOperandSegments]> {
         - slice: an array section can be described with a slice triple,
         - typeparams: for emboxing a derived type with LEN type parameters,
         - accessMap: unused/experimental.
+        - allocator_idx: specify special allocator to use.
   }];
 
   let arguments = (ins
@@ -791,7 +792,8 @@ def fir_EmboxOp : fir_Op<"embox", [NoMemoryEffect, AttrSizedOperandSegments]> {
     Optional<fir_SliceType>:$slice,
     Variadic<AnyIntegerType>:$typeparams,
     Optional<fir_ClassType>:$sourceBox,
-    OptionalAttr<AffineMapAttr>:$accessMap
+    OptionalAttr<AffineMapAttr>:$accessMap,
+    OptionalAttr<I32Attr>:$allocator_idx
   );
 
   let results = (outs BoxOrClassType);
@@ -801,9 +803,11 @@ def fir_EmboxOp : fir_Op<"embox", [NoMemoryEffect, AttrSizedOperandSegments]> {
       "mlir::Value":$memref, CArg<"mlir::Value", "{}">:$shape,
       CArg<"mlir::Value", "{}">:$slice,
       CArg<"mlir::ValueRange", "{}">:$typeparams,
-      CArg<"mlir::Value", "{}">:$sourceBox),
+      CArg<"mlir::Value", "{}">:$sourceBox,
+      CArg<"mlir::IntegerAttr", "{}">:$allocator_idx),
     [{ return build($_builder, $_state, resultTypes, memref, shape, slice,
-                    typeparams, sourceBox, mlir::AffineMapAttr{}); }]>
+                    typeparams, sourceBox, mlir::AffineMapAttr{},
+                    allocator_idx); }]>
   ];
 
   let assemblyFormat = [{
diff --git a/flang/include/flang/Runtime/allocator-registry.h b/flang/include/flang/Runtime/allocator-registry.h
@@ -13,6 +13,8 @@
 #include <cstdlib>
 #include <vector>
 
+static constexpr unsigned kDefaultAllocator = 0;
+
 #define MAX_ALLOCATOR 5
 
 namespace Fortran::runtime {
@@ -37,7 +39,9 @@ struct AllocatorRegistry {
   RT_API_ATTRS constexpr AllocatorRegistry()
       : allocators{{&MallocWrapper, &FreeWrapper}} {}
 #else
-  constexpr AllocatorRegistry() { allocators[0] = {&std::malloc, &std::free}; };
+  constexpr AllocatorRegistry() {
+    allocators[kDefaultAllocator] = {&std::malloc, &std::free};
+  };
 #endif
   RT_API_ATTRS void Register(int, Allocator_t);
   RT_API_ATTRS AllocFct GetAllocator(int pos);
diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -23,6 +23,8 @@
 #include "flang/Optimizer/Support/InternalNames.h"
 #include "flang/Optimizer/Support/TypeCode.h"
 #include "flang/Optimizer/Support/Utils.h"
+#include "flang/Runtime/allocator-registry.h"
+#include "flang/Runtime/descriptor.h"
 #include "flang/Semantics/runtime-type-info.h"
 #include "mlir/Conversion/ArithCommon/AttrToLLVMConverter.h"
 #include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
@@ -1224,8 +1226,8 @@ struct EmboxCommonConversion : public fir::FIROpConversion<OP> {
                                  fir::BaseBoxType boxTy, mlir::Type inputType,
                                  mlir::ConversionPatternRewriter &rewriter,
                                  unsigned rank, mlir::Value eleSize,
-                                 mlir::Value cfiTy,
-                                 mlir::Value typeDesc) const {
+                                 mlir::Value cfiTy, mlir::Value typeDesc,
+                                 int allocatorIdx = kDefaultAllocator) const {
     auto llvmBoxTy = this->lowerTy().convertBoxTypeAsStruct(boxTy, rank);
     bool isUnlimitedPolymorphic = fir::isUnlimitedPolymorphicType(boxTy);
     bool useInputType = fir::isPolymorphicType(boxTy) || isUnlimitedPolymorphic;
@@ -1243,9 +1245,17 @@ struct EmboxCommonConversion : public fir::FIROpConversion<OP> {
                     this->genI32Constant(loc, rewriter, getCFIAttr(boxTy)));
 
     const bool hasAddendum = fir::boxHasAddendum(boxTy);
+
+    // Descriptor used to set the correct value of the extra field.
+    Fortran::runtime::StaticDescriptor<0> staticDescriptor;
+    Fortran::runtime::Descriptor &desc{staticDescriptor.descriptor()};
+    desc.raw().extra = 0;
+    desc.SetAllocIdx(allocatorIdx);
+    if (hasAddendum)
+      desc.SetHasAddendum();
     descriptor =
         insertField(rewriter, loc, descriptor, {kExtraPosInBox},
-                    this->genI32Constant(loc, rewriter, hasAddendum ? 1 : 0));
+                    this->genI32Constant(loc, rewriter, desc.raw().extra));
 
     if (hasAddendum) {
       unsigned typeDescFieldId = getTypeDescFieldId(boxTy);
@@ -1300,6 +1310,13 @@ struct EmboxCommonConversion : public fir::FIROpConversion<OP> {
         typeparams.push_back(substrParams[1]);
     }
 
+    int allocatorIdx = 0;
+    if constexpr (std::is_same_v<BOX, fir::EmboxOp> ||
+                  std::is_same_v<BOX, fir::cg::XEmboxOp>) {
+      if (box.getAllocatorIdx())
+        allocatorIdx = *box.getAllocatorIdx();
+    }
+
     // Write each of the fields with the appropriate values.
     // When emboxing an element to a polymorphic descriptor, use the
     // input type since the destination descriptor type has not the exact
@@ -1321,8 +1338,9 @@ struct EmboxCommonConversion : public fir::FIROpConversion<OP> {
                                     cfiTy.getType(), rewriter, kTypePosInBox);
     }
     auto mod = box->template getParentOfType<mlir::ModuleOp>();
-    mlir::Value descriptor = populateDescriptor(
-        loc, mod, boxTy, inputType, rewriter, rank, eleSize, cfiTy, typeDesc);
+    mlir::Value descriptor =
+        populateDescriptor(loc, mod, boxTy, inputType, rewriter, rank, eleSize,
+                           cfiTy, typeDesc, allocatorIdx);
 
     return {boxTy, descriptor, eleSize};
   }
diff --git a/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp b/flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp
@@ -109,7 +109,7 @@ class EmboxConversion : public mlir::OpRewritePattern<fir::EmboxOp> {
     auto xbox = rewriter.create<fir::cg::XEmboxOp>(
         loc, embox.getType(), embox.getMemref(), shapeOpers, std::nullopt,
         std::nullopt, std::nullopt, std::nullopt, embox.getTypeparams(),
-        embox.getSourceBox());
+        embox.getSourceBox(), embox.getAllocatorIdxAttr());
     LLVM_DEBUG(llvm::dbgs() << "rewriting " << embox << " to " << xbox << '\n');
     rewriter.replaceOp(embox, xbox.getOperation()->getResults());
     return mlir::success();
@@ -145,7 +145,7 @@ class EmboxConversion : public mlir::OpRewritePattern<fir::EmboxOp> {
     auto xbox = rewriter.create<fir::cg::XEmboxOp>(
         loc, embox.getType(), embox.getMemref(), shapeOpers, shiftOpers,
         sliceOpers, subcompOpers, substrOpers, embox.getTypeparams(),
-        embox.getSourceBox());
+        embox.getSourceBox(), embox.getAllocatorIdxAttr());
     LLVM_DEBUG(llvm::dbgs() << "rewriting " << embox << " to " << xbox << '\n');
     rewriter.replaceOp(embox, xbox.getOperation()->getResults());
     return mlir::success();
diff --git a/flang/runtime/allocator-registry.cpp b/flang/runtime/allocator-registry.cpp
@@ -6,7 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "allocator-registry.h"
+#include "flang/Runtime/allocator-registry.h"
 #include "terminator.h"
 
 namespace Fortran::runtime {
diff --git a/flang/runtime/descriptor.cpp b/flang/runtime/descriptor.cpp
@@ -8,13 +8,13 @@
 
 #include "flang/Runtime/descriptor.h"
 #include "ISO_Fortran_util.h"
-#include "allocator-registry.h"
 #include "derived.h"
 #include "memory.h"
 #include "stat.h"
 #include "terminator.h"
 #include "tools.h"
 #include "type-info.h"
+#include "flang/Runtime/allocator-registry.h"
 #include <cassert>
 #include <cstdlib>
 #include <cstring>
diff --git a/flang/test/Fir/embox.fir b/flang/test/Fir/embox.fir
@@ -103,3 +103,27 @@ func.func @fir_dev_issue_1416(%arg0: !fir.ref<!fir.array<40x?xf32>>, %low: index
     fir.call @do_something(%3) : (!fir.box<!fir.array<?xf32>>) -> ()
     return
 }
+
+// CHECK-LABEL: define void @_QPtest_allocator1()
+func.func @_QPtest_allocator1() {
+  %c20 = arith.constant 20 : index
+  %0 = fir.alloca !fir.array<20xi32> {bindc_name = "x", uniq_name = "_QFtest_sliceEx"}
+  %1 = fir.shape %c20 : (index) -> !fir.shape<1>
+  %3 = fir.embox %0(%1) {allocator_idx = 1 : i32} : (!fir.ref<!fir.array<20xi32>>, !fir.shape<1>) -> !fir.box<!fir.array<?xi32>>
+  fir.call @_QPtest_callee(%3) : (!fir.box<!fir.array<?xi32>>) -> ()
+  return
+}
+
+// %{{.*}} = insertvalue { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] } { ptr undef, i64 ptrtoint (ptr getelementptr (i32, ptr null, i32 1) to i64), i32 20240719, i8 1, i8 9, i8 0, i8 2, [1 x [3 x i64]] [[3 x i64] [i64 1, i64 20, i64 ptrtoint (ptr getelementptr (i32, ptr null, i32 1) to i64)]] }
+
+// CHECK-LABEL: define void @_QPtest_allocator2()
+func.func @_QPtest_allocator2() {
+  %c20 = arith.constant 20 : index
+  %0 = fir.alloca !fir.array<20xi32> {bindc_name = "x", uniq_name = "_QFtest_sliceEx"}
+  %1 = fir.shape %c20 : (index) -> !fir.shape<1>
+  %3 = fir.embox %0(%1) {allocator_idx = 3 : i32} : (!fir.ref<!fir.array<20xi32>>, !fir.shape<1>) -> !fir.box<!fir.array<?xi32>>
+  fir.call @_QPtest_callee(%3) : (!fir.box<!fir.array<?xi32>>) -> ()
+  return
+}
+
+// CHECK: %{{.*}} = insertvalue { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] } { ptr undef, i64 ptrtoint (ptr getelementptr (i32, ptr null, i32 1) to i64), i32 20240719, i8 1, i8 9, i8 0, i8 6
