[fir] Add fir.box_rank, fir.box_addr, fir.box_dims and fir.box_elesize conversion

[clementval]

Summary

This patch adds conversion for basic box operations that extract
information from the box.

This patch is part of the upstreaming effort from fir-dev branch.

Co-authored-by: Eric Schweitz <[e-mail]>

[awarzynski]

flang/lib/Optimizer/CodeGen/CodeGen.cpp

32–35

It's great to see some non-magic values used here! Do these numbers correspond to https://github.com/llvm/llvm-project/blob/99ad2079d452f587be050b3867e0ed4856335fb2/flang/lib/Optimizer/CodeGen/TypeConverter.h#L105-L116? If yes, could these be used for the type converter too? Also, could you add a comment here that would explain the origins of these indices? Ta!

flang/test/Fir/convert-to-llvm.fir

646

Shouldn't this be written in a target agnostic way? Like here: https://github.com/llvm/llvm-project/blob/99ad2079d452f587be050b3867e0ed4856335fb2/flang/test/Fir/types-to-llvm.fir#L38 Same comment for other tests here.

653–656

DELETEME?

[clementval]

flang/lib/Optimizer/CodeGen/CodeGen.cpp

32–35

Yes they are the same. Let me do a quick patch to move those value in TypeConverter and use them. I'll base this patch on it.

[clementval]

Address review comments

flang/lib/Optimizer/CodeGen/CodeGen.cpp

32–35

Done in D113553

flang/test/Fir/convert-to-llvm.fir

646

Yeah might be safer.

[awarzynski]

Thanks for the updates,

Makes sense. The remaining comments are mostly nits, so I'm accepting as is.

flang/lib/Optimizer/CodeGen/CodeGen.cpp

59

[nit] It's clear that it's a method. Construct is the key information here.

78

[nit]

116–118

c0 is the offset to get the relevant struct (i.e. to get the base address), and c1 is the offset within the struct (i.e. to get the element address), right?

128

What's cv?

233

DELETEME

This is an implementation detail from getRankFromBox.

[clementval]

Address review comments

flang/lib/Optimizer/CodeGen/CodeGen.cpp

78

This is always 0.

116–118

Yes that's right. c0 is always 0 when accessing the box.

128

constant values.

233

My bad. Was left from the code when there was not constant for the different positions.

[clementval]

Rebase

flang/lib/Optimizer/CodeGen/CodeGen.cpp

 #include "mlir/Pass/Pass.h"
 #include "llvm/ADT/ArrayRef.h"
 
 #define DEBUG_TYPE "flang-codegen"
 
 // fir::LLVMTypeConverter for converting to LLVM IR dialect types.
 #include "TypeConverter.h"
 
+static constexpr unsigned kAddrPosInBox = 0;
+static constexpr unsigned kEleSizePosInBox = 1;
+static constexpr unsigned kRankPosInBox = 3;
+static constexpr unsigned kDimsPosInBox = 7;
+
 namespace {
 /// FIR conversion pattern template
 template <typename FromOp>
 class FIROpConversion : public mlir::ConvertOpToLLVMPattern<FromOp> {
 public:
   explicit FIROpConversion(fir::LLVMTypeConverter &lowering)
       : mlir::ConvertOpToLLVMPattern<FromOp>(lowering) {}
 
 protected:
   mlir::Type convertType(mlir::Type ty) const {
     return lowerTy().convertType(ty);
   }
 
+  mlir::LLVM::ConstantOp
+  genConstantOffset(mlir::Location loc,
+                    mlir::ConversionPatternRewriter &rewriter,
+                    int offset) const {
+    auto ity = lowerTy().offsetType();
+    auto cattr = rewriter.getI32IntegerAttr(offset);
+    return rewriter.create<mlir::LLVM::ConstantOp>(loc, ity, cattr);
+  }
+
+  /// Method to construct code sequence to get the rank from a box.
+  mlir::Value getRankFromBox(mlir::Location loc, mlir::Value box,
+                             mlir::Type resultTy,
+                             mlir::ConversionPatternRewriter &rewriter) const {
+    mlir::LLVM::ConstantOp c0 = genConstantOffset(loc, rewriter, 0);
+    mlir::LLVM::ConstantOp cRank =
+        genConstantOffset(loc, rewriter, kRankPosInBox);
+    auto pty = mlir::LLVM::LLVMPointerType::get(resultTy);
+    auto p = rewriter.create<mlir::LLVM::GEPOp>(
+        loc, pty, mlir::ValueRange{box, c0, cRank});
+    return rewriter.create<mlir::LLVM::LoadOp>(loc, resultTy, p);
+  }
+
+  /// Method to construct code sequence to get the triple for dimension `dim`
+  /// from a box.
+  SmallVector<mlir::Value, 3>
+  getDimsFromBox(mlir::Location loc, ArrayRef<mlir::Type> retTys,
+                 mlir::Value box, mlir::Value dim,
+                 mlir::ConversionPatternRewriter &rewriter) const {
+    mlir::LLVM::ConstantOp c0 = genConstantOffset(loc, rewriter, 0);
+    mlir::LLVM::ConstantOp cDims =
+        genConstantOffset(loc, rewriter, kDimsPosInBox);
+    mlir::LLVM::LoadOp l0 =
+        loadFromOffset(loc, box, c0, cDims, dim, 0, retTys[0], rewriter);
+    mlir::LLVM::LoadOp l1 =
+        loadFromOffset(loc, box, c0, cDims, dim, 1, retTys[1], rewriter);
+    mlir::LLVM::LoadOp l2 =
+        loadFromOffset(loc, box, c0, cDims, dim, 2, retTys[2], rewriter);
+    return {l0.getResult(), l1.getResult(), l2.getResult()};
+  }
+
+  mlir::LLVM::LoadOp
+  loadFromOffset(mlir::Location loc, mlir::Value a, mlir::LLVM::ConstantOp c0,
+                 mlir::LLVM::ConstantOp cDims, mlir::Value dim, int off,
+                 mlir::Type ty,
+                 mlir::ConversionPatternRewriter &rewriter) const {
+    auto pty = mlir::LLVM::LLVMPointerType::get(ty);
+    mlir::LLVM::ConstantOp c = genConstantOffset(loc, rewriter, off);
+    mlir::LLVM::GEPOp p = genGEP(loc, pty, rewriter, a, c0, cDims, dim, c);
+    return rewriter.create<mlir::LLVM::LoadOp>(loc, ty, p);
+  }
+
+  /// Read base address from a fir.box. Returned address has type ty.
+  mlir::Value
+  loadBaseAddrFromBox(mlir::Location loc, mlir::Type ty, mlir::Value box,
+                      mlir::ConversionPatternRewriter &rewriter) const {
+    mlir::LLVM::ConstantOp c0 = genConstantOffset(loc, rewriter, 0);
+    mlir::LLVM::ConstantOp cAddr =
+        genConstantOffset(loc, rewriter, kAddrPosInBox);
+    auto pty = mlir::LLVM::LLVMPointerType::get(ty);
+    mlir::LLVM::GEPOp p = genGEP(loc, pty, rewriter, box, c0, cAddr);
+    return rewriter.create<mlir::LLVM::LoadOp>(loc, ty, p);
+  }
+
+  mlir::Value
+  loadElementSizeFromBox(mlir::Location loc, mlir::Type ty, mlir::Value box,
+                         mlir::ConversionPatternRewriter &rewriter) const {
+    mlir::LLVM::ConstantOp c0 = genConstantOffset(loc, rewriter, 0);
+    mlir::LLVM::ConstantOp c1 =
+        genConstantOffset(loc, rewriter, kEleSizePosInBox);
+    auto pty = mlir::LLVM::LLVMPointerType::get(ty);
+    mlir::LLVM::GEPOp p = genGEP(loc, pty, rewriter, box, c0, c1);
+    return rewriter.create<mlir::LLVM::LoadOp>(loc, ty, p);
+  }
+
+  template <typename... ARGS>
+  mlir::LLVM::GEPOp genGEP(mlir::Location loc, mlir::Type ty,
+                           mlir::ConversionPatternRewriter &rewriter,
+                           mlir::Value base, ARGS... args) const {
+    SmallVector<mlir::Value> cv{args...};
+    return rewriter.create<mlir::LLVM::GEPOp>(loc, ty, base, cv);
+  }
+
   fir::LLVMTypeConverter &lowerTy() const {
     return *static_cast<fir::LLVMTypeConverter *>(this->getTypeConverter());
   }
 };
 
 /// FIR conversion pattern template
 template <typename FromOp>
 class FIROpAndTypeConversion : public FIROpConversion<FromOp> {
matchAndRewrite(fir::AddrOfOp addr, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     auto ty = convertType(addr.getType());
     rewriter.replaceOpWithNewOp<mlir::LLVM::AddressOfOp>(
         addr, ty, addr.symbol().getRootReference().getValue());
     return success();
   }
 };
 
+/// Lower `fir.box_addr` to the sequence of operations to extract the first
+/// element of the box.
+struct BoxAddrOpConversion : public FIROpConversion<fir::BoxAddrOp> {
+  using FIROpConversion::FIROpConversion;
+
+  mlir::LogicalResult
+  matchAndRewrite(fir::BoxAddrOp boxaddr, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    mlir::Value a = adaptor.getOperands()[0];
+    auto loc = boxaddr.getLoc();
+    mlir::Type ty = convertType(boxaddr.getType());
+    if (auto argty = boxaddr.val().getType().dyn_cast<fir::BoxType>()) {
+      rewriter.replaceOp(boxaddr, loadBaseAddrFromBox(loc, ty, a, rewriter));
+    } else {
+      auto c0attr = rewriter.getI32IntegerAttr(0);
+      auto c0 = mlir::ArrayAttr::get(boxaddr.getContext(), c0attr);
+      rewriter.replaceOpWithNewOp<mlir::LLVM::ExtractValueOp>(boxaddr, ty, a,
+                                                              c0);
+    }
+    return success();
+  }
+};
+
+/// Lower `fir.box_dims` to a sequence of operations to extract the requested
+/// dimension infomartion from the boxed value.
+/// Result in a triple set of GEPs and loads.
+struct BoxDimsOpConversion : public FIROpConversion<fir::BoxDimsOp> {
+  using FIROpConversion::FIROpConversion;
+
+  mlir::LogicalResult
+  matchAndRewrite(fir::BoxDimsOp boxdims, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    SmallVector<mlir::Type, 3> resultTypes = {
+        convertType(boxdims.getResult(0).getType()),
+        convertType(boxdims.getResult(1).getType()),
+        convertType(boxdims.getResult(2).getType()),
+    };
+    auto results =
+        getDimsFromBox(boxdims.getLoc(), resultTypes, adaptor.getOperands()[0],
+                       adaptor.getOperands()[1], rewriter);
+    rewriter.replaceOp(boxdims, results);
+    return success();
+  }
+};
+
+/// Lower `fir.box_elesize` to a sequence of operations ro extract the size of
+/// an element in the boxed value.
+struct BoxEleSizeOpConversion : public FIROpConversion<fir::BoxEleSizeOp> {
+  using FIROpConversion::FIROpConversion;
+
+  mlir::LogicalResult
+  matchAndRewrite(fir::BoxEleSizeOp boxelesz, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    mlir::Value a = adaptor.getOperands()[0];
+    auto loc = boxelesz.getLoc();
+    auto ty = convertType(boxelesz.getType());
+    rewriter.replaceOp(boxelesz, loadElementSizeFromBox(loc, ty, a, rewriter));
+    return success();
+  }
+};
+
+/// Lower `fir.box_rank` to the sequence of operation to extract the rank from
+/// the box.
+/// The rank is in 3rd position in the box.
+struct BoxRankOpConversion : public FIROpConversion<fir::BoxRankOp> {
+  using FIROpConversion::FIROpConversion;
+
+  mlir::LogicalResult
+  matchAndRewrite(fir::BoxRankOp boxrank, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    mlir::Value a = adaptor.getOperands()[0];
+    auto loc = boxrank.getLoc();
+    mlir::Type ty = convertType(boxrank.getType());
+    auto result = getRankFromBox(loc, a, ty, rewriter);
+    rewriter.replaceOp(boxrank, result);
+    return success();
+  }
+};
+
 // `fir.call` -> `llvm.call`
 struct CallOpConversion : public FIROpConversion<fir::CallOp> {
   using FIROpConversion::FIROpConversion;
 
   mlir::LogicalResult
   matchAndRewrite(fir::CallOp call, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     SmallVector<mlir::Type> resultTys;
void runOnOperation() override final {
     auto mod = getModule();
     if (!forcedTargetTriple.empty()) {
       fir::setTargetTriple(mod, forcedTargetTriple);
     }
 
     auto *context = getModule().getContext();
     fir::LLVMTypeConverter typeConverter{getModule()};
     mlir::OwningRewritePatternList pattern(context);
-    pattern
-        .insert<AddcOpConversion, AddrOfOpConversion, CallOpConversion,
-                ConvertOpConversion, DivcOpConversion, ExtractValueOpConversion,
+    pattern.insert<AddcOpConversion, AddrOfOpConversion, BoxAddrOpConversion,
+                   BoxDimsOpConversion, BoxEleSizeOpConversion,
+                   BoxRankOpConversion, CallOpConversion, ConvertOpConversion,
+                   DivcOpConversion, ExtractValueOpConversion,
                    HasValueOpConversion, GlobalOpConversion,
                    InsertOnRangeOpConversion, InsertValueOpConversion,
                    NegcOpConversion, MulcOpConversion, SelectOpConversion,
                    SelectRankOpConversion, SubcOpConversion, UndefOpConversion,
                    UnreachableOpConversion, ZeroOpConversion>(typeConverter);
     mlir::populateStdToLLVMConversionPatterns(typeConverter, pattern);
     mlir::arith::populateArithmeticToLLVMConversionPatterns(typeConverter,
                                                             pattern);
     mlir::ConversionTarget target{*context};
     target.addLegalDialect<mlir::LLVM::LLVMDialect>();
 
     // required NOPs for applying a full conversion
     target.addLegalOp<mlir::ModuleOp>();

flang/lib/Optimizer/CodeGen/TypeConverter.h

addConversion([&](mlir::TupleType tuple) {
       llvm::SmallVector<mlir::Type> members;
       for (auto mem : inMembers)
         members.push_back(convertType(mem).cast<mlir::Type>());
       return mlir::LLVM::LLVMStructType::getLiteral(&getContext(), members,
                                                     /*isPacked=*/false);
     });
   }
 
+  // i32 is used here because LLVM wants i32 constants when indexing into struct
+  // types. Indexing into other aggregate types is more flexible.
+  mlir::Type offsetType() { return mlir::IntegerType::get(&getContext(), 32); }
+
   // fir.type<name(p : TY'...){f : TY...}>  -->  llvm<"%name = { ty... }">
   mlir::Type convertRecordType(fir::RecordType derived) {
     auto name = derived.getName();
     auto st = mlir::LLVM::LLVMStructType::getIdentified(&getContext(), name);
     llvm::SmallVector<mlir::Type> members;
     for (auto mem : derived.getTypeList()) {
       members.push_back(convertType(mem.second).cast<mlir::Type>());
     }

flang/test/Fir/convert-to-llvm.fir

 // CHECK:         %[[X:.*]] = llvm.extractvalue %[[ARG0]][0 : i32] : !llvm.struct<(f128, f128)>
 // CHECK:         %[[Y:.*]] = llvm.extractvalue %[[ARG0]][1 : i32] : !llvm.struct<(f128, f128)>
 // CHECK:         %[[CONVERTX:.*]] = llvm.fptrunc %[[X]] : f128 to f16
 // CHECK:         %[[CONVERTY:.*]] = llvm.fptrunc %[[Y]] : f128 to f16
 // CHECK:         %[[STRUCT0:.*]] = llvm.mlir.undef : !llvm.struct<(f16, f16)>
 // CHECK:         %[[STRUCT1:.*]] = llvm.insertvalue %[[CONVERTX]], %[[STRUCT0]][0 : i32] : !llvm.struct<(f16, f16)>
 // CHECK:         %[[STRUCT2:.*]] = llvm.insertvalue %[[CONVERTY]], %[[STRUCT1]][1 : i32] : !llvm.struct<(f16, f16)>
 // CHECK:         llvm.return %[[STRUCT2]] : !llvm.struct<(f16, f16)>
+
+// -----
+
+// Test `fir.box_rank` conversion.
+
+func @extract_rank(%arg0: !fir.box<!fir.array<*:f64>>) -> i32 {
+  %0 = fir.box_rank %arg0 : (!fir.box<!fir.array<*:f64>>) -> i32
+  return %0 : i32
+}
+
+// CHECK-LABEL: llvm.func @extract_rank(
+// CHECK-SAME:                          %[[ARG0:.*]]: !llvm.ptr<struct<(ptr<f64>, i64, i32, i8, i8, i8, i8)>>) -> i32
+// CHECK:         %[[C0:.*]] = llvm.mlir.constant(0 : i32) : i32
+// CHECK:         %[[CRANK:.*]] = llvm.mlir.constant(3 : i32) : i32
+// CHECK:         %[[GEP:.*]] = llvm.getelementptr %[[ARG0]][%[[C0]], %[[CRANK]]] : (!llvm.ptr<struct<(ptr<f64>, i64, i32, i8, i8, i8, i8)>>, i32, i32) -> !llvm.ptr<i32>
+// CHECK:         %[[RANK:.*]] = llvm.load %[[GEP]] : !llvm.ptr<i32>
+// CHECK:         llvm.return %[[RANK]] : i32
+
+func @foo1(%arg0: !fir.box<f64>) -> i32 {
+  %0 = fir.box_rank %arg0 : (!fir.box<f64>) -> i32
+  return %0 : i32
+}
+
+// -----
+
+// Test `fir.box_addr` conversion.
+
+func @extract_addr(%arg0: !fir.box<!fir.array<*:f64>>) -> !fir.ref<f64> {
+  %0 = fir.box_addr %arg0 : (!fir.box<!fir.array<*:f64>>) -> !fir.ref<f64>
+  return %0 : !fir.ref<f64>
+}
+
+// CHECK-LABEL: llvm.func @extract_addr(
+// CHECK-SAME:                          %[[ARG0:.*]]: !llvm.ptr<struct<(ptr<f64>, i64, i32, i8, i8, i8, i8)>>) -> !llvm.ptr<f64>
+// CHECK:         %[[C0:.*]] = llvm.mlir.constant(0 : i32) : i32
+// CHECK:         %[[CADDR:.*]] = llvm.mlir.constant(0 : i32) : i32
+// CHECK:         %[[GEP:.*]] = llvm.getelementptr %[[ARG0]][%[[C0]], %[[CADDR]]] : (!llvm.ptr<struct<(ptr<f64>, i64, i32, i8, i8, i8, i8)>>, i32, i32) -> !llvm.ptr<ptr<f64>>
+// CHECK:         %[[ADDR:.*]] = llvm.load %[[GEP]] : !llvm.ptr<ptr<f64>>
+// CHECK:         llvm.return %[[ADDR]] : !llvm.ptr<f64>
+
+// -----
+
+// Test `fir.box_dims` conversion.
+
+func @extract_dims(%arg0: !fir.box<!fir.array<*:f64>>) -> index {
+  %c1 = arith.constant 0 : i32
+  %0:3 = fir.box_dims %arg0, %c1 : (!fir.box<!fir.array<*:f64>>, i32) -> (index, index, index)
+  return %0 : index
+}
+
+// CHECK-LABEL: llvm.func @extract_dims(
+// CHECK-SAME:                          %[[ARG0:.*]]: !llvm.ptr<struct<(ptr<f64>, i64, i32, i8, i8, i8, i8)>>) -> i64
+// CHECK:         %[[C0_1:.*]] = llvm.mlir.constant(0 : i32) : i32
+// CHECK:         %[[C0:.*]] = llvm.mlir.constant(0 : i32) : i32
+// CHECK:         %[[CDIMS:.*]] = llvm.mlir.constant(7 : i32) : i32
+// CHECK:         %[[C0_2:.*]] = llvm.mlir.constant(0 : i32) : i32
+// CHECK:         %[[GEP0:.*]] = llvm.getelementptr %[[ARG0]][%[[C0]], %[[CDIMS]], %[[C0_1]], %[[C0_2]]] : (!llvm.ptr<struct<(ptr<f64>, i64, i32, i8, i8, i8, i8)>>, i32, i32, i32, i32) -> !llvm.ptr<i64>
+// CHECK:         %[[LOAD0:.*]] = llvm.load %[[GEP0]] : !llvm.ptr<i64>
+// CHECK:         %[[C1:.*]] = llvm.mlir.constant(1 : i32) : i32
+// CHECK:         %[[GEP1:.*]] = llvm.getelementptr %[[ARG0]][%[[C0]], %[[CDIMS]], %[[C0_1]], %[[C1]]] : (!llvm.ptr<struct<(ptr<f64>, i64, i32, i8, i8, i8, i8)>>, i32, i32, i32, i32) -> !llvm.ptr<i64>
+// CHECK:         %[[LOAD1:.*]] = llvm.load %[[GEP1]] : !llvm.ptr<i64>
+// CHECK:         %[[C2:.*]] = llvm.mlir.constant(2 : i32) : i32
+// CHECK:         %[[GEP2:.*]] = llvm.getelementptr %[[ARG0]][%[[C0]], %[[CDIMS]], %[[C0_1]], %[[C2]]] : (!llvm.ptr<struct<(ptr<f64>, i64, i32, i8, i8, i8, i8)>>, i32, i32, i32, i32) -> !llvm.ptr<i64>
+// CHECK:         %[[LOAD2:.*]] = llvm.load %[[GEP2]] : !llvm.ptr<i64>
+// CHECK:         llvm.return %[[LOAD0]] : i64
+
+// -----
+
+// Test `fir.box_elesize` conversion.
+
+func @extract_elesize(%arg0: !fir.box<f32>) -> i32 {
+  %0 = fir.box_elesize %arg0 : (!fir.box<f32>) -> i32
+  return %0 : i32
+}
+
+// CHECK-LABEL: llvm.func @extract_elesize(
+// CHECK-SAME:                             %[[ARG0:.*]]: !llvm.ptr<struct<(ptr<f32>, i64, i32, i8, i8, i8, i8)>>) -> i32
+// CHECK:         %[[C0:.*]] = llvm.mlir.constant(0 : i32) : i32
+// CHECK:         %[[CELESIZE:.*]] = llvm.mlir.constant(1 : i32) : i32
+// CHECK:         %[[GEP:.*]] = llvm.getelementptr %[[ARG0]][%[[C0]], %[[CELESIZE]]] : (!llvm.ptr<struct<(ptr<f32>, i64, i32, i8, i8, i8, i8)>>, i32, i32) -> !llvm.ptr<i32>
+// CHECK:         %[[ELE_SIZE:.*]] = llvm.load %[[GEP]] : !llvm.ptr<i32>
+// CHECK:         llvm.return %[[ELE_SIZE]] : i32