[fir] TargetRewrite: Rewrite fir.address_of(func)

[rovka]

Summary

Rewrite AddrOfOp if taking the address of a function.

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

[clementval]

Thanks for working on that. Just a small comment on my side.

flang/lib/Optimizer/CodeGen/TargetRewrite.cpp

352

flang/lib/Optimizer/CodeGen/TargetRewrite.cpp

void runOnOperation() override final {
     // Convert ops in target-specific patterns.
     mod.walk([&](mlir::Operation *op) {
       if (auto call = dyn_cast<fir::CallOp>(op)) {
         if (!hasPortableSignature(call.getFunctionType()))
           convertCallOp(call);
       } else if (auto dispatch = dyn_cast<DispatchOp>(op)) {
         if (!hasPortableSignature(dispatch.getFunctionType()))
           convertCallOp(dispatch);
+      } else if (auto addr = dyn_cast<AddrOfOp>(op)) {
+        if (addr.getType().isa<mlir::FunctionType>() &&
+            !hasPortableSignature(addr.getType()))
+          convertAddrOp(addr);
       }
     });
 
     clearMembers();
   }
 
   mlir::ModuleOp getModule() { return getOperation(); }
 
public:
   void lowerComplexSignatureArg(A cmplx, B &newInTys) {
     if (noComplexConversion)
       newInTys.push_back(cmplx);
     else
       for (auto &tup : specifics->complexArgumentType(cmplx.getElementType()))
         newInTys.push_back(std::get<mlir::Type>(tup));
   }
 
+  /// Taking the address of a function. Modify the signature as needed.
+  void convertAddrOp(AddrOfOp addrOp) {
+    rewriter->setInsertionPoint(addrOp);
+    auto addrTy = addrOp.getType().cast<mlir::FunctionType>();
+    llvm::SmallVector<mlir::Type> newResTys;
+    llvm::SmallVector<mlir::Type> newInTys;
+    for (mlir::Type ty : addrTy.getResults()) {
+      llvm::TypeSwitch<mlir::Type>(ty)
+          .Case<fir::ComplexType>([&](fir::ComplexType ty) {
+            lowerComplexSignatureRes(ty, newResTys, newInTys);
+          })
+          .Case<mlir::ComplexType>([&](mlir::ComplexType ty) {
+            lowerComplexSignatureRes(ty, newResTys, newInTys);
+          })
+          .Default([&](mlir::Type ty) { newResTys.push_back(ty); });
+    }
+    llvm::SmallVector<mlir::Type> trailingInTys;
+    for (mlir::Type ty : addrTy.getInputs()) {
+      llvm::TypeSwitch<mlir::Type>(ty)
+          .Case<BoxCharType>([&](BoxCharType box) {
+            if (noCharacterConversion) {
+              newInTys.push_back(box);
+            } else {
+              for (auto &tup : specifics->boxcharArgumentType(box.getEleTy())) {
+                auto attr = std::get<CodeGenSpecifics::Attributes>(tup);
+                auto argTy = std::get<mlir::Type>(tup);
+                llvm::SmallVector<mlir::Type> &vec =
+                    attr.isAppend() ? trailingInTys : newInTys;
+                vec.push_back(argTy);
+              }
+            }
+          })
+          .Case<fir::ComplexType>([&](fir::ComplexType ty) {
+            lowerComplexSignatureArg(ty, newInTys);
+          })
+          .Case<mlir::ComplexType>([&](mlir::ComplexType ty) {
+            lowerComplexSignatureArg(ty, newInTys);
+          })
+          .Default([&](mlir::Type ty) { newInTys.push_back(ty); });
+    }
+    // append trailing input types
+    newInTys.insert(newInTys.end(), trailingInTys.begin(), trailingInTys.end());
+    // replace this op with a new one with the updated signature
+    auto newTy = rewriter->getFunctionType(newInTys, newResTys);
+    auto newOp =
+        rewriter->create<AddrOfOp>(addrOp.getLoc(), newTy, addrOp.symbol());
+    replaceOp(addrOp, newOp.getResult());
+  }
+
   /// Convert the type signatures on all the functions present in the module.
   /// As the type signature is being changed, this must also update the
   /// function itself to use any new arguments, etc.
   mlir::LogicalResult convertTypes(mlir::ModuleOp mod) {
     for (auto fn : mod.getOps<mlir::FuncOp>())
       convertSignature(fn);
     return mlir::success();
   }

flang/test/Fir/target-rewrite-boxchar.fir

func @boxchararg() {
   return
 }
 
 fir.global @name constant : !fir.char<1,9> {
   %str = fir.string_lit "Your name"(9) : !fir.char<1,9>
   //constant 1
   fir.has_value %str : !fir.char<1,9>
 }
+
+// Test that we rewrite the fir.address_of operator
+// INT32-LABEL: @addrof
+// INT64-LABEL: @addrof
+func @addrof() {
+  // INT32: {{.*}} = fir.address_of(@boxcharcallee) : (!fir.ref<!fir.char<1,?>>, i32) -> ()
+  // INT64: {{.*}} = fir.address_of(@boxcharcallee) : (!fir.ref<!fir.char<1,?>>, i64) -> ()
+  %f = fir.address_of(@boxcharcallee) : (!fir.boxchar<1>) -> ()
+  return
+}

flang/test/Fir/target-rewrite-complex.fir

func private @mlircomplexf32(%z1: complex<f32>, %z2: complex<f32>) -> complex<f32> {
   // PPC: [[V:%[0-9A-Za-z]+]] = fir.load [[ADDRC]] : !fir.ref<complex<f32>>
   // PPC: [[ADDRT_2:%[0-9A-Za-z]+]] = fir.alloca tuple<f32, f32>
   // PPC: [[ADDRC_2:%[0-9A-Za-z]+]] = fir.convert [[ADDRT_2]] : (!fir.ref<tuple<f32, f32>>) -> !fir.ref<complex<f32>>
   // PPC: fir.store [[V]] to [[ADDRC_2]] : !fir.ref<complex<f32>>
   // PPC: [[RES:%[0-9A-Za-z]+]] = fir.load [[ADDRT_2]] : !fir.ref<tuple<f32, f32>>
   // PPC: return [[RES]] : tuple<f32, f32>
   return %0 : complex<f32>
 }
+
+// Test that we rewrite the fir.address_of operator.
+// I32-LABEL: func @addrof()
+// X64-LABEL: func @addrof()
+// AARCH64-LABEL: func @addrof()
+// PPC-LABEL: func @addrof()
+func @addrof() {
+  // I32: {{%.*}} = fir.address_of(@returncomplex4) : () -> i64
+  // X64: {{%.*}} = fir.address_of(@returncomplex4) : () -> !fir.vector<2:!fir.real<4>>
+  // AARCH64: {{%.*}} = fir.address_of(@returncomplex4) : () -> tuple<!fir.real<4>, !fir.real<4>>
+  // PPC: {{%.*}} = fir.address_of(@returncomplex4) : () -> tuple<!fir.real<4>, !fir.real<4>>
+  %r = fir.address_of(@returncomplex4) : () -> !fir.complex<4>
+
+  // I32: {{%.*}} = fir.address_of(@paramcomplex4) : (!fir.ref<tuple<!fir.real<4>, !fir.real<4>>>) -> ()
+  // X64: {{%.*}} = fir.address_of(@paramcomplex4) : (!fir.vector<2:!fir.real<4>>) -> ()
+  // AARCH64: {{%.*}} = fir.address_of(@paramcomplex4) : (!fir.array<2x!fir.real<4>>) -> ()
+  // PPC: {{%.*}} = fir.address_of(@paramcomplex4) : (!fir.real<4>, !fir.real<4>) -> ()
+  %p = fir.address_of(@paramcomplex4) : (!fir.complex<4>) -> ()
+  return
+}