[flang] Handle emboxing of a fir.ref<none> to an unlimited polymorphic box

[clementval]

Summary

When an array element is extracted from an unlimited polymorphic array, the
emboxing of this element has to retrive the type code and element size from
the initial array. This patch retrive this information through the extracted
type descriptor.

This situation can be found in code like:

subroutine sub1(a)
  class(*) :: a(:)
  select type (x=>a(1))
  type is (integer)
    x = 10
  end select
end subroutine

[jeanPerier]

Makes sense with the current IR. I think might be cleaner to try to improve the indexing/assumed type handling in lowering so that the FIR operation that produces a new fir.class hold the input fir.class as operand and there is no need to find it back.

Do you have the motivating Fortran example for this so that I can try to see if HLFIR addressing might help in the future ?

[clementval]

In D141274#4036051, @jeanPerier wrote:

Makes sense with the current IR. I think might be cleaner to try to improve the indexing/assumed type handling in lowering so that the FIR operation that produces a new fir.class hold the input fir.class as operand and there is no need to find it back.

Do you have the motivating Fortran example for this so that I can try to see if HLFIR addressing might help in the future ?

Agreed it would be nice to improve this and make it cleaner.

Here is an example:

program test
  integer :: arr(2) = [ 10,20 ]
  call sub1(arr)
Contains
  subroutine sub1(arr)
    class(*) :: arr(:)
    select type (x => arr(2))
    type is (integer)
      a = 10
    end select
  end subroutine
end

[PeteSteinfeld]

All builds and tests correctly and looks good.

flang/lib/Optimizer/CodeGen/CodeGen.cpp

consDescriptorPrefix(BOX box, mlir::Type inputType,
     }
 
     // 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
     // information.
     auto [eleSize, cfiTy] = getSizeAndTypeCode(
         loc, rewriter, useInputType ? inputType : boxTy.getEleTy(), typeparams);
+
+    // When emboxing a fir.ref<none> to an unlimited polymorphic box, get the
+    // type code and element size from the box used to extract the type desc.
+    if (fir::isUnlimitedPolymorphicType(boxTy) &&
+        inputType.isa<mlir::NoneType>() && typeDesc) {
+      if (auto *typeDescOp = typeDesc.getDefiningOp()) {
+        if (auto loadOp = mlir::dyn_cast<mlir::LLVM::LoadOp>(typeDescOp)) {
+          if (auto *gepOp = loadOp.getAddr().getDefiningOp()) {
+            if (auto gep = mlir::dyn_cast<mlir::LLVM::GEPOp>(gepOp)) {
+              mlir::Type idxTy = this->lowerTy().indexType();
+              eleSize = this->loadElementSizeFromBox(loc, idxTy, gep.getBase(),
+                                                     rewriter);
+              cfiTy = this->getValueFromBox(loc, gep.getBase(), cfiTy.getType(),
+                                            rewriter, kTypePosInBox);
+            }
+          }
+        }
+      }
+    }
+
     auto mod = box->template getParentOfType<mlir::ModuleOp>();
     mlir::Value descriptor = populateDescriptor(
         loc, mod, boxTy, inputType, rewriter, rank, eleSize, cfiTy, typeDesc);
 
     return {boxTy, descriptor, eleSize};
   }
 
   std::tuple<fir::BaseBoxType, mlir::Value, mlir::Value>

flang/test/Fir/polymorphic.fir

 }
 
 // CHECK-LABEL: @_QMpolymorphic_testPtest_embox()
 // CHECK: %[[ALLOCA_DESC:.*]] = alloca { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]], ptr, [1 x i64] }
 // CHECK: store { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]], ptr, [1 x i64] } { ptr @_QFEy, i64 ptrtoint (ptr getelementptr (i32, ptr null, i32 1) to i64), i32 20180515, i8 1, i8 9, {{.*}}, ptr %[[ALLOCA_DESC]]
 // CHECK: %[[LOADED_DESC:.*]] = load { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]], ptr, [1 x i64] }, ptr %[[ALLOCA_DESC]], align 8
 // CHECK: store { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]], ptr, [1 x i64] } %[[LOADED_DESC]], ptr @_QFEx, align 8
 
+// Test emboxing of an array element from an unlimited polymorphic array.
+
+func.func @_QMunlimitedPsub1(%arg0: !fir.class<!fir.array<?xnone>> {fir.bindc_name = "a"}) {
+  %c1_i64 = arith.constant 1 : i64
+  %c1_i64_0 = arith.constant 1 : i64
+  %0 = arith.subi %c1_i64, %c1_i64_0 : i64
+  %1 = fir.coordinate_of %arg0, %0 : (!fir.class<!fir.array<?xnone>>, i64) -> !fir.ref<none>
+  %2 = fir.box_tdesc %arg0 : (!fir.class<!fir.array<?xnone>>) -> !fir.tdesc<none>
+  %3 = fir.embox %1 tdesc %2 : (!fir.ref<none>, !fir.tdesc<none>) -> !fir.class<none>
+  fir.select_type %3 : !fir.class<none> [#fir.type_is<i32>, ^bb1, unit, ^bb2]
+^bb1:
+  %4 = fir.box_addr %3 : (!fir.class<none>) -> !fir.ref<i32>
+  %c10_i32 = arith.constant 10 : i32
+  fir.store %c10_i32 to %4 : !fir.ref<i32>
+  cf.br ^bb2
+^bb2:  // 2 preds: ^bb0, ^bb1
+  return
+}
+
+// CHECK-LABEL: define void @_QMunlimitedPsub1(
+// CHECK-SAME: ptr %[[ARRAY:.*]]) {
+// CHECK: %[[BOX:.*]] = alloca { ptr, i64, i32, i8, i8, i8, i8, ptr, [1 x i64] }
+// CHECK: %[[TYPE_DESC_GEP:.*]] = getelementptr { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]], ptr, [1 x i64] }, ptr %[[ARRAY]], i32 0, i32 8
+// CHECK: %[[TYPE_DESC:.*]] = load ptr, ptr %[[TYPE_DESC_GEP]]
+// CHECK: %[[ELE_SIZE_GEP:.*]] = getelementptr { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]], ptr, [1 x i64] }, ptr %0, i32 0, i32 1
+// CHECK: %[[ELE_SIZE:.*]] = load i64, ptr %[[ELE_SIZE_GEP]]
+// CHECK: %[[TYPE_CODE_GEP:.*]] = getelementptr { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]], ptr, [1 x i64] }, ptr %0, i32 0, i32 4
+// CHECK: %[[TYPE_CODE:.*]] = load i32, ptr %[[TYPE_CODE_GEP]]
+// CHECK: %{{.*}} = insertvalue { ptr, i64, i32, i8, i8, i8, i8, ptr, [1 x i64] } undef, i64 %[[ELE_SIZE]], 1
+// CHECK: %[[TYPE_CODE_TRUNC:.*]] = trunc i32 %[[TYPE_CODE]] to i8
+// CHECK: %{{.*}} = insertvalue { ptr, i64, i32, i8, i8, i8, i8, ptr, [1 x i64] } %{{.*}}, i8 %[[TYPE_CODE_TRUNC]], 4
+// CHECK: %{{.*}} = insertvalue { ptr, i64, i32, i8, i8, i8, i8, ptr, [1 x i64] } %{{.*}}, ptr %[[TYPE_DESC]], 7
+// CHECK: store { ptr, i64, i32, i8, i8, i8, i8, ptr, [1 x i64] } %24, ptr %[[BOX]]
+
 
 fir.global internal @_QFEx : !fir.class<!fir.ptr<!fir.array<?xnone>>> {
   %0 = fir.zero_bits !fir.ptr<!fir.array<?xnone>>
   %c0 = arith.constant 0 : index
   %1 = fir.shape %c0 : (index) -> !fir.shape<1>
   %2 = fir.embox %0(%1) : (!fir.ptr<!fir.array<?xnone>>, !fir.shape<1>) -> !fir.class<!fir.ptr<!fir.array<?xnone>>>
   fir.has_value %2 : !fir.class<!fir.ptr<!fir.array<?xnone>>>
 }
 
 fir.global internal @_QFEy target : !fir.array<1xi32> {
   %0 = fir.undefined !fir.array<1xi32>
   fir.has_value %0 : !fir.array<1xi32>
 }
 
 func.func private @_FortranAioBeginExternalListOutput(i32, !fir.ref<i8>, i32) -> !fir.ref<i8> attributes {fir.io, fir.runtime}
 func.func private @_FortranAioOutputDescriptor(!fir.ref<i8>, !fir.box<none>) -> i1 attributes {fir.io, fir.runtime}
 func.func private @_FortranAioEndIoStatement(!fir.ref<i8>) -> i32 attributes {fir.io, fir.runtime}
 fir.global linkonce @_QQcl.2E2F64756D6D792E66393000 constant : !fir.char<1,12> {
   %0 = fir.string_lit "./dummy.f90\00"(12) : !fir.char<1,12>
   fir.has_value %0 : !fir.char<1,12>
 }
+