[flang][codegen] Add a conversion for !fir.coordinate_of - part 1

This patch extends the `FIRToLLVMLowering` pass in Flang by adding a hook to transform `!fir.coordinate_of` into a sequence of LLVM MLIR instructions. The following cases are currently supported: 1. the input object is a `!fir.complex` (wrapped in e.g. `!fir.ref` or `!fir.box`) 2. the input object is wrapped in a `!fir.box` (including e.g. `!fir.array`). Note that `!fir.complex` inside a `!fir.box` falls under case 1. above (i.e. it's a special case regardless of the wrapping type). This is part of the upstreaming effort from the `!fir-dev` branch in [1]. Differential Revision: https://reviews.llvm.org/D114159 Co-authored-by: Jean Perier <jperier@nvidia.com> Co-authored-by: Eric Schweitz <eschweitz@nvidia.com> Co-authored-by: V Donaldson <vdonaldson@nvidia.com>
llvm · Dec 10, 2021 · 75db341 · 75db341
1 parent cf58b9c
commit 75db341
Show file tree

Hide file tree

Showing 7 changed files with 438 additions and 1 deletion.
diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -2910,6 +2910,168 @@ struct FieldIndexOpConversion : public FIROpConversion<fir::FieldIndexOp> {
   }
 };
 
+/// Convert to (memory) reference to a reference to a subobject.
+/// The coordinate_of op is a Swiss army knife operation that can be used on
+/// (memory) references to records, arrays, complex, etc. as well as boxes.
+/// With unboxed arrays, there is the restriction that the array have a static
+/// shape in all but the last column.
+struct CoordinateOpConversion
+    : public FIROpAndTypeConversion<fir::CoordinateOp> {
+  using FIROpAndTypeConversion::FIROpAndTypeConversion;
+
+  mlir::LogicalResult
+  doRewrite(fir::CoordinateOp coor, mlir::Type ty, OpAdaptor adaptor,
+            mlir::ConversionPatternRewriter &rewriter) const override {
+    mlir::ValueRange operands = adaptor.getOperands();
+
+    mlir::Location loc = coor.getLoc();
+    mlir::Value base = operands[0];
+    mlir::Type baseObjectTy = coor.getBaseType();
+    mlir::Type objectTy = fir::dyn_cast_ptrOrBoxEleTy(baseObjectTy);
+    assert(objectTy && "fir.coordinate_of expects a reference type");
+
+    // Complex type - basically, extract the real or imaginary part
+    if (fir::isa_complex(objectTy)) {
+      mlir::LLVM::ConstantOp c0 =
+          genConstantIndex(loc, lowerTy().indexType(), rewriter, 0);
+      SmallVector<mlir::Value> offs = {c0, operands[1]};
+      mlir::Value gep = genGEP(loc, ty, rewriter, base, offs);
+      rewriter.replaceOp(coor, gep);
+      return success();
+    }
+
+    // Box type - get the base pointer from the box
+    if (auto boxTy = baseObjectTy.dyn_cast<fir::BoxType>()) {
+      doRewriteBox(coor, ty, operands, loc, rewriter);
+      return success();
+    }
+
+    // Sequence type (e.g. fir.array)
+    if (auto arrTy = objectTy.dyn_cast<fir::SequenceType>()) {
+      doRewriteSequence(loc);
+      return success();
+    }
+
+    return rewriter.notifyMatchFailure(
+        coor, "fir.coordinate_of base operand has unsupported type");
+  }
+
+  unsigned getFieldNumber(fir::RecordType ty, mlir::Value op) const {
+    return fir::hasDynamicSize(ty)
+               ? op.getDefiningOp()
+                     ->getAttrOfType<mlir::IntegerAttr>("field")
+                     .getInt()
+               : getIntValue(op);
+  }
+
+  int64_t getIntValue(mlir::Value val) const {
+    assert(val && val.dyn_cast<mlir::OpResult>() && "must not be null value");
+    mlir::Operation *defop = val.getDefiningOp();
+
+    if (auto constOp = dyn_cast<mlir::arith::ConstantIntOp>(defop))
+      return constOp.value();
+    if (auto llConstOp = dyn_cast<mlir::LLVM::ConstantOp>(defop))
+      if (auto attr = llConstOp.value().dyn_cast<mlir::IntegerAttr>())
+        return attr.getValue().getSExtValue();
+    fir::emitFatalError(val.getLoc(), "must be a constant");
+  }
+
+private:
+  void doRewriteBox(fir::CoordinateOp coor, mlir::Type ty,
+                    mlir::ValueRange operands, mlir::Location loc,
+                    mlir::ConversionPatternRewriter &rewriter) const {
+    mlir::Type boxObjTy = coor.getBaseType();
+    assert(boxObjTy.dyn_cast<fir::BoxType>() && "This is not a `fir.box`");
+
+    mlir::Value boxBaseAddr = operands[0];
+
+    // 1. SPECIAL CASE (uses `fir.len_param_index`):
+    //   %box = ... : !fir.box<!fir.type<derived{len1:i32}>>
+    //   %lenp = fir.len_param_index len1, !fir.type<derived{len1:i32}>
+    //   %addr = coordinate_of %box, %lenp
+    if (coor.getNumOperands() == 2) {
+      mlir::Operation *coordinateDef = (*coor.coor().begin()).getDefiningOp();
+      if (isa_and_nonnull<fir::LenParamIndexOp>(coordinateDef)) {
+        TODO(loc,
+             "fir.coordinate_of - fir.len_param_index is not supported yet");
+      }
+    }
+
+    // 2. GENERAL CASE:
+    // 2.1. (`fir.array`)
+    //   %box = ... : !fix.box<!fir.array<?xU>>
+    //   %idx = ... : index
+    //   %resultAddr = coordinate_of %box, %idx : !fir.ref<U>
+    // 2.2 (`fir.derived`)
+    //   %box = ... : !fix.box<!fir.type<derived_type{field_1:i32}>>
+    //   %idx = ... : i32
+    //   %resultAddr = coordinate_of %box, %idx : !fir.ref<i32>
+    // 2.3 (`fir.derived` inside `fir.array`)
+    //   %box = ... : !fir.box<!fir.array<10 x !fir.type<derived_1{field_1:f32, field_2:f32}>>>
+    //   %idx1 = ... : index
+    //   %idx2 = ... : i32
+    //   %resultAddr = coordinate_of %box, %idx1, %idx2 : !fir.ref<f32>
+    // 2.4. TODO: Either document or disable any other case that the following
+    //  implementation might convert.
+    mlir::LLVM::ConstantOp c0 =
+        genConstantIndex(loc, lowerTy().indexType(), rewriter, 0);
+    mlir::Value resultAddr =
+        loadBaseAddrFromBox(loc, getBaseAddrTypeFromBox(boxBaseAddr.getType()),
+                            boxBaseAddr, rewriter);
+    auto currentObjTy = fir::dyn_cast_ptrOrBoxEleTy(boxObjTy);
+    mlir::Type voidPtrTy = ::getVoidPtrType(coor.getContext());
+
+    for (unsigned i = 1, last = operands.size(); i < last; ++i) {
+      if (auto arrTy = currentObjTy.dyn_cast<fir::SequenceType>()) {
+        if (i != 1)
+          TODO(loc, "fir.array nested inside other array and/or derived type");
+        // Applies byte strides from the box. Ignore lower bound from box
+        // since fir.coordinate_of indexes are zero based. Lowering takes care
+        // of lower bound aspects. This both accounts for dynamically sized
+        // types and non contiguous arrays.
+        auto idxTy = lowerTy().indexType();
+        mlir::Value off = genConstantIndex(loc, idxTy, rewriter, 0);
+        for (unsigned index = i, lastIndex = i + arrTy.getDimension();
+             index < lastIndex; ++index) {
+          mlir::Value stride =
+              loadStrideFromBox(loc, operands[0], index - i, rewriter);
+          auto sc = rewriter.create<mlir::LLVM::MulOp>(loc, idxTy,
+                                                       operands[index], stride);
+          off = rewriter.create<mlir::LLVM::AddOp>(loc, idxTy, sc, off);
+        }
+        auto voidPtrBase =
+            rewriter.create<mlir::LLVM::BitcastOp>(loc, voidPtrTy, resultAddr);
+        SmallVector<mlir::Value> args{voidPtrBase, off};
+        resultAddr = rewriter.create<mlir::LLVM::GEPOp>(loc, voidPtrTy, args);
+        i += arrTy.getDimension() - 1;
+        currentObjTy = arrTy.getEleTy();
+      } else if (auto recTy = currentObjTy.dyn_cast<fir::RecordType>()) {
+        auto recRefTy =
+            mlir::LLVM::LLVMPointerType::get(lowerTy().convertType(recTy));
+        mlir::Value nxtOpnd = operands[i];
+        auto memObj =
+            rewriter.create<mlir::LLVM::BitcastOp>(loc, recRefTy, resultAddr);
+        llvm::SmallVector<mlir::Value> args = {memObj, c0, nxtOpnd};
+        currentObjTy = recTy.getType(getFieldNumber(recTy, nxtOpnd));
+        auto llvmCurrentObjTy = lowerTy().convertType(currentObjTy);
+        auto gep = rewriter.create<mlir::LLVM::GEPOp>(
+            loc, mlir::LLVM::LLVMPointerType::get(llvmCurrentObjTy), args);
+        resultAddr =
+            rewriter.create<mlir::LLVM::BitcastOp>(loc, voidPtrTy, gep);
+      } else {
+        fir::emitFatalError(loc, "unexpected type in coordinate_of");
+      }
+    }
+
+    rewriter.replaceOpWithNewOp<mlir::LLVM::BitcastOp>(coor, ty, resultAddr);
+    return;
+  }
+
+  void doRewriteSequence(mlir::Location loc) const {
+    TODO(loc, "fir.coordinate_of codegen for sequence types");
+  }
+};
+
 } // namespace
 
 namespace {
@@ -2939,7 +3101,7 @@ class FIRToLLVMLowering : public fir::FIRToLLVMLoweringBase<FIRToLLVMLowering> {
         BoxIsAllocOpConversion, BoxIsArrayOpConversion, BoxIsPtrOpConversion,
         BoxProcHostOpConversion, BoxRankOpConversion, BoxTypeDescOpConversion,
         CallOpConversion, CmpcOpConversion, ConstcOpConversion,
-        ConvertOpConversion, DispatchOpConversion, DispatchTableOpConversion,
+        ConvertOpConversion, CoordinateOpConversion, DispatchOpConversion, DispatchTableOpConversion,
         DTEntryOpConversion, DivcOpConversion, EmboxOpConversion,
         EmboxCharOpConversion, EmboxProcOpConversion, ExtractValueOpConversion,
         FieldIndexOpConversion, FirEndOpConversion, FreeMemOpConversion,

diff --git a/flang/lib/Optimizer/CodeGen/TypeConverter.h b/flang/lib/Optimizer/CodeGen/TypeConverter.h
@@ -31,6 +31,7 @@ static constexpr unsigned kTypePosInBox = 4;
 static constexpr unsigned kAttributePosInBox = 5;
 static constexpr unsigned kF18AddendumPosInBox = 6;
 static constexpr unsigned kDimsPosInBox = 7;
+static constexpr unsigned kStridePosInDim = 2;
 static constexpr unsigned kOptTypePtrPosInBox = 8;
 static constexpr unsigned kOptRowTypePosInBox = 9;
 // Position of the different values in [dims]

diff --git a/flang/test/Fir/Todo/cordinate_of_1.fir b/flang/test/Fir/Todo/cordinate_of_1.fir
@@ -0,0 +1,12 @@
+// RUN: %not_todo_cmd fir-opt --fir-to-llvm-ir="target=x86_64-unknown-linux-gnu" %s 2>&1 | FileCheck %s
+
+// `fir.coordinate_of` - derived type with `fir.len_param_index`. As
+// `fir.len_param_index` is not implemented yet, that's the error that's
+// currently being generated (this error is generated before trying to convert
+// `fir.coordinate_of`)
+func @coordinate_box_derived_with_fir_len(%arg0: !fir.box<!fir.type<derived_2{len1:i32}>>) {
+// CHECK: not yet implemented fir.len_param_index codegen
+  %e = fir.len_param_index len1, !fir.type<derived_2{len1:i32}>
+  %q = fir.coordinate_of %arg0, %e : (!fir.box<!fir.type<derived_2{len1:i32}>>, !fir.len) -> !fir.ref<i32>
+  return
+}
diff --git a/flang/test/Fir/Todo/cordinate_of_2.fir b/flang/test/Fir/Todo/cordinate_of_2.fir
@@ -0,0 +1,10 @@
+// RUN: %not_todo_cmd fir-opt --fir-to-llvm-ir="target=x86_64-unknown-linux-gnu" %s 2>&1 | FileCheck %s
+
+// CHECK: not yet implemented fir.array nested inside other array and/or derived type
+
+// `!fir.coordinate_of` - `!fir.array` inside "boxed" `!fir.type`
+func @coordinate_box_array_inside_derived(%arg0: !fir.box<!fir.type<derived_2{field_1:!fir.array<10 x i32>, field_2:i32}>>, %arg1 : index) {
+   %idx0 = arith.constant 0 : i32
+   %q = fir.coordinate_of %arg0, %idx0, %arg1 : (!fir.box<!fir.type<derived_2{field_1:!fir.array<10 x i32>, field_2:i32}>>, i32, index) -> !fir.ref<f32>
+   return
+}
diff --git a/flang/test/Fir/Todo/cordinate_of_3.fir b/flang/test/Fir/Todo/cordinate_of_3.fir
@@ -0,0 +1,10 @@
+// RUN: %not_todo_cmd fir-opt --fir-to-llvm-ir="target=x86_64-unknown-linux-gnu" %s 2>&1 | FileCheck %s
+
+// CHECK: not yet implemented fir.array nested inside other array and/or derived type
+
+// `fir.coordinate_of` - `fir.array` inside "boxed" `!fir.type<derived_1{!fir.type<derived_2{}>}` (i.e. nested `!fir.type`)
+func @coordinate_box_array_inside_derived(%arg0: !fir.box<!fir.type<derived_1{field_1:!fir.type<derived_2{field_2:!fir.array<10 x i32>}>}>>, %arg1 : index) {
+   %idx0 = arith.constant 0 : i32
+   %q = fir.coordinate_of %arg0, %idx0, %idx0, %arg1 : (!fir.box<!fir.type<derived_1{field_1:!fir.type<derived_2{field_2:!fir.array<10 x i32>}>}>>, i32, i32, index) -> !fir.ref<f32>
+   return
+}
diff --git a/flang/test/Fir/Todo/cordinate_of_4.fir b/flang/test/Fir/Todo/cordinate_of_4.fir
@@ -0,0 +1,11 @@
+// RUN: %not_todo_cmd fir-opt --fir-to-llvm-ir="target=x86_64-unknown-linux-gnu" %s 2>&1 | FileCheck %s
+
+// `!fir.coordinate_of` - derived type with `!fir.len_param_index`. As
+// `!fir.len_param_index` is not implemented yet, the error that we hit is
+// related to `!fir.len_param_index` rather than `!fir.coordinate_of`.
+func @coordinate_box_derived_with_fir_len(%arg0: !fir.box<!fir.type<derived_2{len1:i32}>>) {
+// CHECK: not yet implemented fir.len_param_index codegen
+  %e = fir.len_param_index len1, !fir.type<derived_2{len1:i32}>
+  %q = fir.coordinate_of %arg0, %e : (!fir.box<!fir.type<derived_2{len1:i32}>>, !fir.len) -> !fir.ref<i32>
+  return
+}