[flang][fir] Added load-store elimination to fir.store

[shraiysh]

Summary

This patch adds canonicalization for fir.store - the store operation is
eliminated when it is right next to a load from the same address.

[kiranchandramohan]

Please wait for a review from the FIR team.

I will just make a note in passing that there is a fir-memref-datflow-opt that does some load to store forwarding and store removal in flang/lib/Optimizer/Transforms/MemRefDataFlowOpt.cpp. But that pass is probably not handling the cases here.

[shraiysh]

Please wait for a review from the FIR team.

Sure. It would be great if you could add some relevant reviewers as I am not aware who should review this except the already added ones.

there is a fir-memref-datflow-opt that does some load to store forwarding and store removal in flang/lib/Optimizer/Transforms/MemRefDataFlowOpt.cpp. But that pass is probably not handling the cases here.

I do not know about them. I will check out what they do, but with the current HEAD, these stores were not getting eliminated - that is correct.

[shraiysh]

In D126643#3546004, @shraiysh wrote:

Please wait for a review from the FIR team.

Sure. It would be great if you could add some relevant reviewers as I am not aware who should review this except the already added ones.

Apologies, just saw that you already did! :)

[schweitz]

Kiran makes a good point that there is a strong desire to perform a true mem-to-reg translation on FIR which would eliminate dead stores and register promote memory traffic to stack slots, etc.

It's not clear to me that there is a lot of IR that is simply the back-to-back Op pattern of:

%1 = fir.load %addr
fir.store %1 to %addr

flang/lib/Optimizer/Dialect/FIROps.cpp

3035

This seems like an extremely narrow application of dead-store removal. Does it actually fire all that much?

[shraiysh]

there is a strong desire to perform a true mem-to-reg translation on FIR

Now that I think about it, this elimination is too special and it would be covered more effectively in a mem2reg kind of pass. Would make more sense to improve that instead of adding it separately. I will abandon this change and maybe work on that later.

flang/include/flang/Optimizer/Dialect/FIROps.td

def fir_StoreOp : fir_Op<"store", [MemoryEffects<[MemWrite]>]> {
   }];
 
   let arguments = (ins AnyType:$value,
                    Arg<AnyReferenceLike, "", [MemWrite]>:$memref);
 
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 
+  let hasCanonicalizeMethod = 1;
+
   let extraClassDeclaration = [{
     static mlir::Type elementType(mlir::Type refType);
   }];
 }
 
 def fir_SaveResultOp : fir_Op<"save_result", [AttrSizedOperandSegments]> {
   let summary = [{
     save an array, box, or record function result SSA-value to a memory location

flang/lib/Optimizer/Dialect/FIROps.cpp

 //===----------------------------------------------------------------------===//
 // StoreOp
 //===----------------------------------------------------------------------===//
 
 mlir::Type fir::StoreOp::elementType(mlir::Type refType) {
   return fir::dyn_cast_ptrEleTy(refType);
 }
 
+mlir::LogicalResult
+fir::StoreOp::canonicalize(fir::StoreOp op, mlir::PatternRewriter &rewriter) {
+  mlir::Operation *definingOp = op.getValue().getDefiningOp();
+  if (!definingOp)
+    return mlir::failure();
+  fir::LoadOp loadOp = mlir::dyn_cast<fir::LoadOp>(definingOp);
+
+  // Load-store elimination
+  if (loadOp && op.getMemref() == loadOp.getMemref() &&
+      loadOp->getNextNode() == op) {
+    rewriter.eraseOp(op);
+    return mlir::success();
+  }
+  return mlir::failure();
+}
+
 mlir::ParseResult fir::StoreOp::parse(mlir::OpAsmParser &parser,
                                       mlir::OperationState &result) {
   mlir::Type type;
   mlir::OpAsmParser::UnresolvedOperand oper;
   mlir::OpAsmParser::UnresolvedOperand store;
   if (parser.parseOperand(oper) || parser.parseKeyword("to") ||
       parser.parseOperand(store) ||
       parser.parseOptionalAttrDict(result.attributes) ||

flang/test/Fir/fir-canonicalize.fir

+// RUN: fir-opt --split-input-file --canonicalize %s | FileCheck %s
+
+// This file will be used for canonicalizations that are defined as part of
+// the operation (not external patterns)
+
+// CHECK-LABEL: @canonicalize_store_simple
+func.func @canonicalize_store_simple(%arg0: !fir.ref<i32>) {
+  %0 = fir.load %arg0 : !fir.ref<i32>
+  fir.store %0 to %arg0 : !fir.ref<i32> // should be eliminated
+  return
+}
+
+// CHECK-NOT: fir.store
+
+// -----
+
+// CHECK-LABEL: @canonicalize_store_argument
+func.func @canonicalize_store_argument(%arg0: !fir.ref<i32>, %arg1: i32) {
+  %0 = fir.load %arg0 : !fir.ref<i32>
+  fir.store %arg1 to %arg0 : !fir.ref<i32> // should not be eliminated
+  return
+}
+
+// CHECK-NOT: fir.load
+// CHECK: fir.store
+
+// -----
+
+func.func private @foo(%arg0: i32)
+
+// CHECK-LABEL: @canonicalize_store_load_multiple_uses
+func.func @canonicalize_store_load_multiple_uses(%arg0: !fir.ref<i32>, %arg1: i32) {
+  %0 = fir.load %arg0 : !fir.ref<i32>
+  fir.store %0 to %arg0 : !fir.ref<i32> // should be eliminated
+  call @foo(%0) : (i32) -> ()
+  return
+}
+
+// CHECK: fir.load
+// CHECK-NOT: fir.store
+// CHECK: call @foo
+
+// -----
+
+func.func private @foo(!fir.ref<i32>)
+
+// CHECK-LABEL: @canonicalize_store_use_before_store
+func.func @canonicalize_store_use_before_store(%arg0: !fir.ref<i32>, %1: i32) {
+  %0 = fir.load %arg0 : !fir.ref<i32>
+  fir.store %1 to %arg0 : !fir.ref<i32> // should not be eliminated
+  call @foo(%arg0) : (!fir.ref<i32>) -> () // called with new value
+  fir.store %0 to %arg0 : !fir.ref<i32> // should not be eliminated
+  call @foo(%arg0) : (!fir.ref<i32>) -> () // called with old value
+  return
+}
+
+// CHECK: fir.store
+// CHECK: call @foo
+// CHECK: fir.store
+// CHECK: call @foo