[flang][CodeGen] Transform `fir.{store|load}` to `llvm.{store|load}`

[awarzynski]

Summary

This patch extends the FIRToLLVMLowering pass in Flang by adding a
hook to transform fir.store/fir.load to llvm.store/fir.load,
respectively.

This is part of the upstreaming effort from the fir-dev branch in [1].

[1] https://github.com/flang-compiler/f18-llvm-project

Patch originally written by:
Co-authored-by: Eric Schweitz <[e-mail]>
Co-authored-by: Jean Perier <[e-mail]>

[awarzynski]

I wanted to add a test for fir.ref<fir.box>, but that does not work yet (it works fine on fir-dev). I wasn't sure why, so my comment in "flang/test/Fir/convert-to-llvm.fir" is a bit generic. Any pointers? Thanks!

[clementval]

In D113090#3105376, @awarzynski wrote:

I wanted to add a test for fir.ref<fir.box>, but that does not work yet (it works fine on fir-dev). I wasn't sure why, so my comment in "flang/test/Fir/convert-to-llvm.fir" is a bit generic. Any pointers? Thanks!

Not all of the type conversions have been upstreamed. I usually upstream the one I need with the op conversion I'm upstreaming.

flang/lib/Optimizer/CodeGen/TypeConverter.h

[clementval]

Test is failing on the pre-check.

[clementval]

flang/lib/Optimizer/CodeGen/CodeGen.cpp

326

Can you address this formatting issue?

[kiranchandramohan]

In D113090#3105389, @clementval wrote:

In D113090#3105376, @awarzynski wrote:

I wanted to add a test for fir.ref<fir.box>, but that does not work yet (it works fine on fir-dev). I wasn't sure why, so my comment in "flang/test/Fir/convert-to-llvm.fir" is a bit generic. Any pointers? Thanks!

Not all of the type conversions have been upstreamed. I usually upstream the one I need with the op conversion I'm upstreaming.

flang/lib/Optimizer/CodeGen/TypeConverter.h

So does @awarzynski have to bring in the convertBoxType function?

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

179

missed the parameters.
Would it be good to match the parameters and the store operands?

200

missed the parameters.
Would it be good to match the parameters and the load operands?

[clementval]

In D113090#3105615, @kiranchandramohan wrote:

In D113090#3105389, @clementval wrote:

In D113090#3105376, @awarzynski wrote:

I wanted to add a test for fir.ref<fir.box>, but that does not work yet (it works fine on fir-dev). I wasn't sure why, so my comment in "flang/test/Fir/convert-to-llvm.fir" is a bit generic. Any pointers? Thanks!

Not all of the type conversions have been upstreamed. I usually upstream the one I need with the op conversion I'm upstreaming.

flang/lib/Optimizer/CodeGen/TypeConverter.h

So does @awarzynski have to bring in the convertBoxType function?

Or it will come with D112961 if anyone review it.

[clementval]

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

179

+1

[mehdi_amini]

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

182

In general: please use CHECK-LABEL for delimiting the test (that helps FileCheck, in particular the print in case of failure is much friendlier) and minimize the checks, while capturing what matters (here I suspect types and argument flowing in the right place it what you're checking (I think that's also what Kiran hinted at):

// CHECK-LABEL: llvm.func @test_store
// CHECK-SAME: %[[ARG0:.*]]: <type here>, %[[ARG1:.*]]: <type here>
// CHECK:    llvm.store %[[ARG0]], .....

[awarzynski]

In D113090#3105647, @clementval wrote:

In D113090#3105615, @kiranchandramohan wrote:

So does @awarzynski have to bring in the convertBoxType function?

Or it will come with D112961 if anyone review it.

Tl;Dr Perhaps we should just add more tests later?

Longer version - I've tried applying D112961, but that didn't work. It depends on 2 other patches, so perhaps I wasn't applying it correctly? I've just tried my usual arc patch D112961. Regardless, if we prefer to have all tests added in one patch then I can set the dependencies accordingly and re-visit arc patch.

[awarzynski]

flang/lib/Optimizer/CodeGen/CodeGen.cpp

326

Sorry about that :/

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

182

haha, as I was in the process of changing that :-)

[awarzynski]

Address PR comments (refine tests, fix formatting)

[clementval]

LGTM

[awarzynski]

Rebased, addes tests for fir.box. This has already been accepted by @clementval, so I will merge this if there are no new comments.

[clementval]

Still LGTM

flang/lib/Optimizer/CodeGen/CodeGen.cpp

void selectMatchAndRewrite(fir::LLVMTypeConverter &lowering, OP select,
 
   llvm::SmallVector<mlir::Block *> destinations;
   llvm::SmallVector<mlir::ValueRange> destinationsOperands;
   mlir::Block *defaultDestination;
   mlir::ValueRange defaultOperands;
   llvm::SmallVector<int32_t> caseValues;
 
   for (unsigned t = 0; t != conds; ++t) {
     mlir::Block *dest = select.getSuccessor(t);
     auto destOps = select.getSuccessorOperands(adaptor.getOperands(), t);
     const mlir::Attribute &attr = cases[t];
     if (auto intAttr = attr.template dyn_cast<mlir::IntegerAttr>()) {
       destinations.push_back(dest);
       destinationsOperands.push_back(destOps.hasValue() ? *destOps
                                                         : ValueRange());
       caseValues.push_back(intAttr.getInt());
       continue;
struct SelectOpConversion : public FIROpConversion<fir::SelectOp> {
   mlir::LogicalResult
   matchAndRewrite(fir::SelectOp op, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     selectMatchAndRewrite<fir::SelectOp>(lowerTy(), op, adaptor, rewriter);
     return success();
   }
 };
 
+/// `fir.load` --> `llvm.load`
+struct LoadOpConversion : public FIROpConversion<fir::LoadOp> {
+  using FIROpConversion::FIROpConversion;
+
+  mlir::LogicalResult
+  matchAndRewrite(fir::LoadOp load, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    // fir.box is a special case because it is considered as an ssa values in
+    // fir, but it is lowered as a pointer to a descriptor. So fir.ref<fir.box>
+    // and fir.box end up being the same llvm types and loading a
+    // fir.ref<fir.box> is actually a no op in LLVM.
+    if (load.getType().isa<fir::BoxType>()) {
+      rewriter.replaceOp(load, adaptor.getOperands()[0]);
+    } else {
+      mlir::Type ty = convertType(load.getType());
+      ArrayRef<NamedAttribute> at = load->getAttrs();
+      rewriter.replaceOpWithNewOp<mlir::LLVM::LoadOp>(
+          load, ty, adaptor.getOperands(), at);
+    }
+    return success();
+  }
+};
+
 /// conversion of fir::SelectRankOp to an if-then-else ladder
 struct SelectRankOpConversion : public FIROpConversion<fir::SelectRankOp> {
   using FIROpConversion::FIROpConversion;
 
   mlir::LogicalResult
   matchAndRewrite(fir::SelectRankOp op, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     selectMatchAndRewrite<fir::SelectRankOp>(lowerTy(), op, adaptor, rewriter);
     return success();
   }
 };
 
-// convert to LLVM IR dialect `undef`
+/// `fir.store` --> `llvm.store`
+struct StoreOpConversion : public FIROpConversion<fir::StoreOp> {
+  using FIROpConversion::FIROpConversion;
+
+  mlir::LogicalResult
+  matchAndRewrite(fir::StoreOp store, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    if (store.value().getType().isa<fir::BoxType>()) {
+      // fir.box value is actually in memory, load it first before storing it.
+      mlir::Location loc = store.getLoc();
+      mlir::Type boxPtrTy = adaptor.getOperands()[0].getType();
+      auto val = rewriter.create<mlir::LLVM::LoadOp>(
+          loc, boxPtrTy.cast<mlir::LLVM::LLVMPointerType>().getElementType(),
+          adaptor.getOperands()[0]);
+      rewriter.replaceOpWithNewOp<mlir::LLVM::StoreOp>(
+          store, val, adaptor.getOperands()[1]);
+    } else {
+      rewriter.replaceOpWithNewOp<mlir::LLVM::StoreOp>(
+          store, adaptor.getOperands()[0], adaptor.getOperands()[1]);
+    }
+    return success();
+  }
+};
+
+/// convert to LLVM IR dialect `undef`
 struct UndefOpConversion : public FIROpConversion<fir::UndefOp> {
   using FIROpConversion::FIROpConversion;
 
   mlir::LogicalResult
   matchAndRewrite(fir::UndefOp undef, OpAdaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<mlir::LLVM::UndefOp>(
         undef, convertType(undef.getType()));
     return success();
   }
 };
 
-// convert to LLVM IR dialect `unreachable`
+/// `fir.unreachable` --> `llvm.unreachable`
 struct UnreachableOpConversion : public FIROpConversion<fir::UnreachableOp> {
   using FIROpConversion::FIROpConversion;
 
   mlir::LogicalResult
   matchAndRewrite(fir::UnreachableOp unreach, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<mlir::LLVM::UnreachableOp>(unreach);
     return success();
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,
-                HasValueOpConversion, GlobalOpConversion,
-                InsertOnRangeOpConversion, InsertValueOpConversion,
-                NegcOpConversion, MulcOpConversion, SelectOpConversion,
-                SelectRankOpConversion, SubcOpConversion, UndefOpConversion,
+    pattern.insert<AddcOpConversion, AddrOfOpConversion, CallOpConversion,
+                   ConvertOpConversion, DivcOpConversion,
+                   ExtractValueOpConversion, HasValueOpConversion,
+                   GlobalOpConversion, InsertOnRangeOpConversion,
+                   InsertValueOpConversion, LoadOpConversion, NegcOpConversion,
+                   MulcOpConversion, SelectOpConversion, SelectRankOpConversion,
+                   StoreOpConversion, 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/test/Fir/convert-to-llvm.fir

 func @test_unreachable() {
   fir.unreachable
 }
 
 // -----
 
 // Test `fir.select` operation conversion pattern.
 // Check that the if-then-else ladder is correctly constructed and that we
 // branch to the correct block.
 
 func @select(%arg : index, %arg2 : i32) -> i32 {
   %0 = arith.constant 1 : i32
   %1 = arith.constant 2 : i32
   %2 = arith.constant 3 : i32
   %3 = arith.constant 4 : i32
   fir.select %arg:index [ 1, ^bb1(%0:i32),
                           2, ^bb2(%2,%arg,%arg2:i32,index,i32),
                           3, ^bb3(%arg2,%2:i32,i32),
                           4, ^bb4(%1:i32),
                           unit, ^bb5 ]
   ^bb1(%a : i32) :
     return %a : i32
   ^bb2(%b : i32, %b2 : index, %b3:i32) :
     %castidx = arith.index_cast %b2 : index to i32
     %4 = arith.addi %b, %castidx : i32
     %5 = arith.addi %4, %b3 : i32
     return %5 : i32
   ^bb3(%c:i32, %c2:i32) :
     %6 = arith.addi %c, %c2 : i32
     return %6 : i32
   ^bb4(%d : i32) :
     return %d : i32
   ^bb5 :
     %zero = arith.constant 0 : i32
     return %zero : i32
 }
 
 // CHECK-LABEL: func @select(

 // 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.store` --> `llvm.store` conversion
+
+func @test_store_index(%val_to_store : index, %addr : !fir.ref<index>) {
+  fir.store %val_to_store to %addr : !fir.ref<index>
+  return
+}
+
+// CHECK-LABEL:   llvm.func @test_store_index
+// CHECK-SAME:    (%[[arg0:.*]]: i64, %[[arg1:.*]]: !llvm.ptr<i64>) {
+// CHECK-NEXT:    llvm.store %[[arg0]], %[[arg1]] : !llvm.ptr<i64>
+// CHECK-NEXT:    llvm.return
+// CHECK-NEXT:  }
+
+func @test_store_box(%array : !fir.ref<!fir.box<!fir.array<?x?xf32>>>, %box : !fir.box<!fir.array<?x?xf32>>) {
+  fir.store %box to %array : !fir.ref<!fir.box<!fir.array<?x?xf32>>>
+  return
+}
+
+// CHECK-LABEL:  llvm.func @test_store_box
+// CHECK-SAME:  (%[[arg0:.*]]: !llvm.ptr<struct<(ptr<f{{.*}}>, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<2 x array<3 x i{{.*}}>>)>>,
+// CHECK-SAME:  %[[arg1:.*]]: !llvm.ptr<struct<(ptr<f{{.*}}>, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<2 x array<3 x i{{.*}}>>)>>) {
+// CHECK-NEXT:  %[[box_to_store:.*]] = llvm.load %arg1 : !llvm.ptr<struct<(ptr<f{{.*}}>, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<2 x array<3 x i{{.*}}>>)>>
+// CHECK-NEXT:  llvm.store %[[box_to_store]], %[[arg0]] : !llvm.ptr<struct<(ptr<f{{.*}}>, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<2 x array<3 x i{{.*}}>>)>>
+// CHECK-NEXT:  llvm.return
+// CHECK-NEXT:  }
+
+// -----
+
+// Test `fir.load` --> `llvm.load` conversion
+
+func @test_load_index(%addr : !fir.ref<index>) {
+  %0 = fir.load %addr : !fir.ref<index>
+  return
+}
+
+// CHECK-LABEL: llvm.func @test_load_index(
+// CHECK-SAME:  %[[arg1:.*]]: !llvm.ptr<i64>) {
+// CHECK-NEXT:    %0 = llvm.load %[[arg1]] : !llvm.ptr<i64>
+// CHECK-NEXT:    llvm.return
+// CHECK-NEXT:  }
+
+func @test_load_box(%addr : !fir.ref<!fir.box<!fir.array<10xf32>>>) {
+  %0 = fir.load %addr : !fir.ref<!fir.box<!fir.array<10xf32>>>
+  return
+}
+
+// Loading a `fir.ref<!fir.box>> is a no-op
+// CHECK-LABEL: llvm.func @test_load_box
+// CHECK-SAME: (%{{.*}}: !llvm.ptr<struct<(ptr<array<10 x f{{.*}}>>, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, i{{.*}}, array<1 x array<3 x i{{.*}}>>)>>) {
+// CHECK-NEXT:  llvm.return
+// CHECK-NEXT:  }