[flang] Adapt target rewrite for fir.dispatch operation

[clementval]

Summary

Handle rewriting dispatch operation with complex arguments or
return.

sret will be done in a separate patch.

[jeanPerier]

LGTM

[PeteSteinfeld]

All builds and tests correctly and looks good to me.

flang/lib/Optimizer/CodeGen/TargetRewrite.cpp

void convertCallOp(A callOp) {
     // to call.
     int dropFront = 0;
     if constexpr (std::is_same_v<std::decay_t<A>, fir::CallOp>) {
       if (!callOp.getCallee()) {
         newInTys.push_back(fnTy.getInput(0));
         newOpers.push_back(callOp.getOperand(0));
         dropFront = 1;
       }
+    } else {
+      dropFront = 1; // First operand is the polymorphic object.
     }
 
     // Determine the rewrite function, `wrap`, for the result value.
     llvm::Optional<std::function<mlir::Value(mlir::Operation *)>> wrap;
     if (fnTy.getResults().size() == 1) {
       mlir::Type ty = fnTy.getResult(0);
       llvm::TypeSwitch<mlir::Type>(ty)
           .template Case<fir::ComplexType>([&](fir::ComplexType cmplx) {
             wrap = rewriteCallComplexResultType(loc, cmplx, newResTys, newInTys,
                                                 newOpers);
           })
           .template Case<mlir::ComplexType>([&](mlir::ComplexType cmplx) {
             wrap = rewriteCallComplexResultType(loc, cmplx, newResTys, newInTys,
                                                 newOpers);
           })
           .Default([&](mlir::Type ty) { newResTys.push_back(ty); });
     } else if (fnTy.getResults().size() > 1) {
       TODO(loc, "multiple results not supported yet");
     }
 
     llvm::SmallVector<mlir::Type> trailingInTys;
     llvm::SmallVector<mlir::Value> trailingOpers;
+    unsigned passArgShift = 0;
     for (auto e : llvm::enumerate(
              llvm::zip(fnTy.getInputs().drop_front(dropFront),
                        callOp.getOperands().drop_front(dropFront)))) {
       mlir::Type ty = std::get<0>(e.value());
       mlir::Value oper = std::get<1>(e.value());
       unsigned index = e.index();
       llvm::TypeSwitch<mlir::Type>(ty)
           .template Case<fir::BoxCharType>([&](fir::BoxCharType boxTy) {
for (auto e : llvm::enumerate(
               trailingInTys.push_back(lenType);
               trailingOpers.push_back(len);
             } else {
               newInTys.push_back(tuple);
               newOpers.push_back(oper);
             }
           })
           .Default([&](mlir::Type ty) {
+            if constexpr (std::is_same_v<std::decay_t<A>, fir::DispatchOp>) {
+              if (callOp.getPassArgPos() && *callOp.getPassArgPos() == index)
+                passArgShift = newOpers.size() - *callOp.getPassArgPos();
+            }
             newInTys.push_back(ty);
             newOpers.push_back(oper);
           });
     }
     newInTys.insert(newInTys.end(), trailingInTys.begin(), trailingInTys.end());
     newOpers.insert(newOpers.end(), trailingOpers.begin(), trailingOpers.end());
     if constexpr (std::is_same_v<std::decay_t<A>, fir::CallOp>) {
       fir::CallOp newCall;
       if (callOp.getCallee()) {
         newCall =
             rewriter->create<A>(loc, *callOp.getCallee(), newResTys, newOpers);
       } else {
         // Force new type on the input operand.
         newOpers[0].setType(mlir::FunctionType::get(
             callOp.getContext(),
             mlir::TypeRange{newInTys}.drop_front(dropFront), newResTys));
         newCall = rewriter->create<A>(loc, newResTys, newOpers);
       }
       LLVM_DEBUG(llvm::dbgs() << "replacing call with " << newCall << '\n');
       if (wrap)
         replaceOp(callOp, (*wrap)(newCall.getOperation()));
       else
         replaceOp(callOp, newCall.getResults());
     } else {
-      // A is fir::DispatchOp
-      TODO(loc, "dispatch not implemented");
+      fir::DispatchOp dispatchOp = rewriter->create<A>(
+          loc, newResTys, rewriter->getStringAttr(callOp.getMethod()),
+          callOp.getOperands()[0], newOpers,
+          rewriter->getI32IntegerAttr(*callOp.getPassArgPos() + passArgShift));
+      if (wrap)
+        replaceOp(callOp, (*wrap)(dispatchOp.getOperation()));
+      else
+        replaceOp(callOp, dispatchOp.getResults());
     }
   }
 
   // Result type fixup for fir::ComplexType and mlir::ComplexType
   template <typename A, typename B>
   void lowerComplexSignatureRes(mlir::Location loc, A cmplx, B &newResTys,
                                 B &newInTys) {
     if (noComplexConversion) {

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

 // I32-LABEL: func private @paramcomplex4(!fir.ref<tuple<!fir.real<4>, !fir.real<4>>> {llvm.align = 4 : i32, llvm.byval = tuple<!fir.real<4>, !fir.real<4>>})
 // X64-LABEL: func private @paramcomplex4(!fir.vector<2:!fir.real<4>>)
 // AARCH64-LABEL: func private @paramcomplex4(!fir.array<2x!fir.real<4>>)
 // PPC-LABEL: func private @paramcomplex4(!fir.real<4>, !fir.real<4>)
 // SPARCV9-LABEL: func private @paramcomplex4(!fir.real<4>, !fir.real<4>)
 func.func private @paramcomplex4(!fir.complex<4>) -> ()
 
 // Test that we rewrite calls to functions that return or accept complex<4>.
-// I32-LABEL: func @callcomplex4()
-// X64-LABEL: func @callcomplex4()
-// AARCH64-LABEL: func @callcomplex4()
-// PPC-LABEL: func @callcomplex4()
-// SPARCV9-LABEL: func @callcomplex4()
-func.func @callcomplex4() {
+// I32-LABEL: func @callcomplex4
+// X64-LABEL: func @callcomplex4
+// AARCH64-LABEL: func @callcomplex4
+// PPC-LABEL: func @callcomplex4
+// SPARCV9-LABEL: func @callcomplex4
+func.func @callcomplex4(%arg0 : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) {
 
   // I32: [[RES:%[0-9A-Za-z]+]] = fir.call @returncomplex4() : () -> i64
   // X64: [[RES:%[0-9A-Za-z]+]] = fir.call @returncomplex4() : () -> !fir.vector<2:!fir.real<4>>
   // AARCH64: [[RES:%[0-9A-Za-z]+]] = fir.call @returncomplex4() : () -> tuple<!fir.real<4>, !fir.real<4>>
   // PPC: [[RES:%[0-9A-Za-z]+]] = fir.call @returncomplex4() : () -> tuple<!fir.real<4>, !fir.real<4>>
   // SPARCV9: [[RES:%[0-9A-Za-z]+]] = fir.call @returncomplex4() : () -> tuple<!fir.real<4>, !fir.real<4>>
   %1 = fir.call @returncomplex4() : () -> !fir.complex<4>
 
func.func @callcomplex4(%arg0 : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) {
   // SPARCV9: [[ADDRT:%[0-9A-Za-z]+]] = fir.alloca tuple<!fir.real<4>, !fir.real<4>>
   // SPARCV9: fir.store [[RES]] to [[ADDRT]] : !fir.ref<tuple<!fir.real<4>, !fir.real<4>>>
   // SPARCV9: [[ADDRC:%[0-9A-Za-z]+]] = fir.convert [[ADDRT]] : (!fir.ref<tuple<!fir.real<4>, !fir.real<4>>>) -> !fir.ref<!fir.complex<4>>
   // SPARCV9: [[V:%[0-9A-Za-z]+]] = fir.load [[ADDRC]] : !fir.ref<!fir.complex<4>>
   // SPARCV9: [[A:%[0-9A-Za-z]+]] = fir.extract_value [[V]], [0 : i32] : (!fir.complex<4>) -> !fir.real<4>
   // SPARCV9: [[B:%[0-9A-Za-z]+]] = fir.extract_value [[V]], [1 : i32] : (!fir.complex<4>) -> !fir.real<4>
   // SPARCV9: fir.call @paramcomplex4([[A]], [[B]]) : (!fir.real<4>, !fir.real<4>) -> ()
   fir.call @paramcomplex4(%1) : (!fir.complex<4>) -> ()
+
+  // I32: [[RES:%[0-9A-Za-z]+]] = fir.dispatch "ret_complex"(%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) (%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) -> i64 {pass_arg_pos = 0 : i32}
+  // X64: [[RES:%[0-9A-Za-z]+]] = fir.dispatch "ret_complex"(%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) (%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) -> !fir.vector<2:!fir.real<4>> {pass_arg_pos = 0 : i32}
+  // AARCH64: [[RES:%[0-9A-Za-z]+]] = fir.dispatch "ret_complex"(%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) (%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) -> tuple<!fir.real<4>, !fir.real<4>> {pass_arg_pos = 0 : i32}
+  // PPC: [[RES:%[0-9A-Za-z]+]] = fir.dispatch "ret_complex"(%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) (%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) -> tuple<!fir.real<4>, !fir.real<4>> {pass_arg_pos = 0 : i32}
+  // SPARCV9: [[RES:%[0-9A-Za-z]+]] = fir.dispatch "ret_complex"(%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) (%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) -> tuple<!fir.real<4>, !fir.real<4>> {pass_arg_pos = 0 : i32}
+  %2 = fir.dispatch "ret_complex"(%arg0 : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) (%arg0 : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) -> !fir.complex<4> {pass_arg_pos = 0 : i32}
+
+  // I32: [[ADDRI64:%[0-9A-Za-z]+]] = fir.alloca i64
+  // I32: fir.store [[RES]] to [[ADDRI64]] : !fir.ref<i64>
+  // I32: [[ADDRC:%[0-9A-Za-z]+]] = fir.convert [[ADDRI64]] : (!fir.ref<i64>) -> !fir.ref<!fir.complex<4>>
+  // I32: [[C:%[0-9A-Za-z]+]] = fir.load [[ADDRC]] : !fir.ref<!fir.complex<4>>
+  // I32: [[ADDRC2:%[0-9A-Za-z]+]] = fir.alloca !fir.complex<4>
+  // I32: fir.store [[C]] to [[ADDRC2]] : !fir.ref<!fir.complex<4>>
+  // I32: [[T:%[0-9A-Za-z]+]] = fir.convert [[ADDRC2]] : (!fir.ref<!fir.complex<4>>) -> !fir.ref<tuple<!fir.real<4>, !fir.real<4>>>
+  // I32: fir.dispatch "with_complex"(%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) (%{{.*}}, [[T]] : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>, !fir.ref<tuple<!fir.real<4>, !fir.real<4>>>) {pass_arg_pos = 0 : i32}
+
+  // X64: [[ADDRV:%[0-9A-Za-z]+]] = fir.alloca !fir.vector<2:!fir.real<4>>
+  // X64: fir.store [[RES]] to [[ADDRV]] : !fir.ref<!fir.vector<2:!fir.real<4>>>
+  // X64: [[ADDRC:%[0-9A-Za-z]+]] = fir.convert [[ADDRV]] : (!fir.ref<!fir.vector<2:!fir.real<4>>>) -> !fir.ref<!fir.complex<4>>
+  // X64: [[V:%[0-9A-Za-z]+]] = fir.load [[ADDRC]] : !fir.ref<!fir.complex<4>>
+  // X64: [[ADDRV2:%[0-9A-Za-z]+]] = fir.alloca !fir.vector<2:!fir.real<4>>
+  // X64: [[ADDRC2:%[0-9A-Za-z]+]] = fir.convert [[ADDRV2]] : (!fir.ref<!fir.vector<2:!fir.real<4>>>) -> !fir.ref<!fir.complex<4>>
+  // X64: fir.store [[V]] to [[ADDRC2]] : !fir.ref<!fir.complex<4>>
+  // X64: [[VRELOADED:%[0-9A-Za-z]+]] = fir.load [[ADDRV2]] : !fir.ref<!fir.vector<2:!fir.real<4>>>
+  // X64: fir.dispatch "with_complex"(%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) (%{{.*}}, [[VRELOADED]] : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>, !fir.vector<2:!fir.real<4>>) {pass_arg_pos = 0 : i32}
+
+  // AARCH64: [[ADDRT:%[0-9A-Za-z]+]] = fir.alloca tuple<!fir.real<4>, !fir.real<4>>
+  // AARCH64: fir.store [[RES]] to [[ADDRT]] : !fir.ref<tuple<!fir.real<4>, !fir.real<4>>>
+  // AARCH64: [[ADDRC:%[0-9A-Za-z]+]] = fir.convert [[ADDRT]] : (!fir.ref<tuple<!fir.real<4>, !fir.real<4>>>) -> !fir.ref<!fir.complex<4>>
+  // AARCH64: [[V:%[0-9A-Za-z]+]] = fir.load [[ADDRC]] : !fir.ref<!fir.complex<4>>
+  // AARCH64: [[ADDRARR:%[0-9A-Za-z]+]] = fir.alloca !fir.array<2x!fir.real<4>>
+  // AARCH64: [[ADDRC2:%[0-9A-Za-z]+]] = fir.convert [[ADDRARR]] : (!fir.ref<!fir.array<2x!fir.real<4>>>) -> !fir.ref<!fir.complex<4>>
+  // AARCH64: fir.store [[V]] to [[ADDRC2]] : !fir.ref<!fir.complex<4>>
+  // AARCH64: [[ARR:%[0-9A-Za-z]+]] = fir.load [[ADDRARR]] : !fir.ref<!fir.array<2x!fir.real<4>>>
+  // AARCH64: fir.dispatch "with_complex"(%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) (%{{.*}}, [[ARR]] : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>, !fir.array<2x!fir.real<4>>) {pass_arg_pos = 0 : i32}
+
+  // PPC: [[ADDRT:%[0-9A-Za-z]+]] = fir.alloca tuple<!fir.real<4>, !fir.real<4>>
+  // PPC: fir.store [[RES]] to [[ADDRT]] : !fir.ref<tuple<!fir.real<4>, !fir.real<4>>>
+  // PPC: [[ADDRC:%[0-9A-Za-z]+]] = fir.convert [[ADDRT]] : (!fir.ref<tuple<!fir.real<4>, !fir.real<4>>>) -> !fir.ref<!fir.complex<4>>
+  // PPC: [[V:%[0-9A-Za-z]+]] = fir.load [[ADDRC]] : !fir.ref<!fir.complex<4>>
+  // PPC: [[A:%[0-9A-Za-z]+]] = fir.extract_value [[V]], [0 : i32] : (!fir.complex<4>) -> !fir.real<4>
+  // PPC: [[B:%[0-9A-Za-z]+]] = fir.extract_value [[V]], [1 : i32] : (!fir.complex<4>) -> !fir.real<4>
+  // PPC: fir.dispatch "with_complex"(%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) (%{{.*}}, [[A]], [[B]] : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>, !fir.real<4>, !fir.real<4>) {pass_arg_pos = 0 : i32}
+
+  // SPARCV9: [[ADDRT:%[0-9A-Za-z]+]] = fir.alloca tuple<!fir.real<4>, !fir.real<4>>
+  // SPARCV9: fir.store [[RES]] to [[ADDRT]] : !fir.ref<tuple<!fir.real<4>, !fir.real<4>>>
+  // SPARCV9: [[ADDRC:%[0-9A-Za-z]+]] = fir.convert [[ADDRT]] : (!fir.ref<tuple<!fir.real<4>, !fir.real<4>>>) -> !fir.ref<!fir.complex<4>>
+  // SPARCV9: [[V:%[0-9A-Za-z]+]] = fir.load [[ADDRC]] : !fir.ref<!fir.complex<4>>
+  // SPARCV9: [[A:%[0-9A-Za-z]+]] = fir.extract_value [[V]], [0 : i32] : (!fir.complex<4>) -> !fir.real<4>
+  // SPARCV9: [[B:%[0-9A-Za-z]+]] = fir.extract_value [[V]], [1 : i32] : (!fir.complex<4>) -> !fir.real<4>
+  // SPARCV9: fir.dispatch "with_complex"(%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) (%{{.*}}, [[A]], [[B]] : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>, !fir.real<4>, !fir.real<4>) {pass_arg_pos = 0 : i32}
+
+  fir.dispatch "with_complex"(%arg0 : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) (%arg0, %2 : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>, !fir.complex<4>) {pass_arg_pos = 0 : i32}
+
+
+  // I32: fir.dispatch "with_complex2"(%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) (%{{.*}}, %{{.*}} : !fir.ref<tuple<!fir.real<4>, !fir.real<4>>>, !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) {pass_arg_pos = 1 : i32}
+  // X64: fir.dispatch "with_complex2"(%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) (%{{.*}}, %{{.*}} : !fir.vector<2:!fir.real<4>>, !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) {pass_arg_pos = 1 : i32}
+  // AARCH64: fir.dispatch "with_complex2"(%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) (%{{.*}}, %{{.*}} : !fir.array<2x!fir.real<4>>, !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) {pass_arg_pos = 1 : i32}
+  // PPC: fir.dispatch "with_complex2"(%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) (%{{.*}}, %{{.*}}, %{{.*}} : !fir.real<4>, !fir.real<4>, !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) {pass_arg_pos = 2 : i32}
+  // SPARCV9: fir.dispatch "with_complex2"(%{{.*}} : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) (%{{.*}}, %{{.*}}, %{{.*}} : !fir.real<4>, !fir.real<4>, !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) {pass_arg_pos = 2 : i32}
+  fir.dispatch "with_complex2"(%arg0 : !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) (%2, %arg0 : !fir.complex<4>, !fir.class<!fir.type<_QMpolymorphic_testTp1{a:i32,b:i32}>>) {pass_arg_pos = 1 : i32}
+  
   return
 }
 
 // Test that we rewrite the signature of a function that accepts a complex<8>.
 // I32-LABEL: func private @paramcomplex8(!fir.ref<tuple<!fir.real<8>, !fir.real<8>>> {llvm.align = 4 : i32, llvm.byval = tuple<!fir.real<8>, !fir.real<8>>})
 // X64-LABEL: func private @paramcomplex8(!fir.real<8>, !fir.real<8>)
 // AARCH64-LABEL: func private @paramcomplex8(!fir.array<2x!fir.real<8>>)
 // PPC-LABEL: func private @paramcomplex8(!fir.real<8>, !fir.real<8>)