[OpenMP] Add TargetAMDGPU support for Complex argument and return types

[TIFitis]

This patch adds TargetAMDGPU offloading support for Complex argument and return types.

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

@llvm/pr-subscribers-flang-codegen

Author: Akash Banerjee (TIFitis)

Changes

This patch adds TargetAMDGPU offloading support for Complex argument and return types.

---

Full diff: https://github.com/llvm/llvm-project/pull/144924.diff

1 Files Affected:

• (modified) flang/lib/Optimizer/CodeGen/Target.cpp (+23-2)

diff --git a/flang/lib/Optimizer/CodeGen/Target.cpp b/flang/lib/Optimizer/CodeGen/Target.cpp
index 7dbf21ce0c125..b60a72e4340b9 100644
--- a/flang/lib/Optimizer/CodeGen/Target.cpp
+++ b/flang/lib/Optimizer/CodeGen/Target.cpp
@@ -1443,14 +1443,35 @@ struct TargetAMDGPU : public GenericTarget<TargetAMDGPU> {
   CodeGenSpecifics::Marshalling
   complexArgumentType(mlir::Location loc, mlir::Type eleTy) const override {
     CodeGenSpecifics::Marshalling marshal;
-    TODO(loc, "handle complex argument types");
+    const auto *sem = &floatToSemantics(kindMap, eleTy);
+    if (sem == &llvm::APFloat::IEEEsingle()) {
+      // Lower COMPLEX(KIND=4) as an array of two element values.
+      marshal.emplace_back(fir::SequenceType::get({2}, eleTy), AT{});
+    } else if (sem == &llvm::APFloat::IEEEdouble()) {
+      // Pass COMPLEX(KIND=8) as two separate arguments.
+      marshal.emplace_back(eleTy, AT{});
+      marshal.emplace_back(eleTy, AT{});
+    } else {
+      typeTodo(sem, loc, "argument");
+    }
     return marshal;
   }
 
   CodeGenSpecifics::Marshalling
   complexReturnType(mlir::Location loc, mlir::Type eleTy) const override {
     CodeGenSpecifics::Marshalling marshal;
-    TODO(loc, "handle complex return types");
+    const auto *sem = &floatToSemantics(kindMap, eleTy);
+    if (sem == &llvm::APFloat::IEEEsingle()) {
+      // Return COMPLEX(KIND=4) as an array of two elements.
+      marshal.emplace_back(fir::SequenceType::get({2}, eleTy), AT{});
+    } else if (sem == &llvm::APFloat::IEEEdouble()) {
+      // Return COMPLEX(KIND=8) via an aggregate with two fields.
+      marshal.emplace_back(mlir::TupleType::get(eleTy.getContext(),
+                                                mlir::TypeRange{eleTy, eleTy}),
+                           AT{});
+    } else {
+      typeTodo(sem, loc, "return");
+    }
     return marshal;
   }
 };

[Copilot]

Pull Request Overview

This PR introduces offloading support for complex argument and return types for TargetAMDGPU. The key changes include:

• Adding marshalling logic in complexArgumentType to handle COMPLEX(KIND=4) and COMPLEX(KIND=8).
• Implementing similar logic in complexReturnType with requirements differing in return type aggregation.

[Copilot]

[nitpick] Consider refactoring the common logic used to compute float semantics in both complexArgumentType and complexReturnType to reduce duplication and improve maintainability.

[TIFitis]

Polite reminder for review, thanks.

[mjklemm]

LGTM