[flang-commits] [flang] [llvm] [Flang][OpenMP] Add global address space to globals for target device (PR #119585)

[Sergio Afonso via flang-commits]

================
@@ -134,16 +134,65 @@ addLLVMOpBundleAttrs(mlir::ConversionPatternRewriter &rewriter,
 }
 
 namespace {
+
+// Creates an existing operation with an AddressOfOp or an AddrSpaceCastOp
+// depending on the existing address spaces of the type.
+mlir::Value createAddrOfOrASCast(mlir::ConversionPatternRewriter &rewriter,
+                                 mlir::Location loc, std::uint64_t globalAS,
+                                 std::uint64_t programAS,
+                                 llvm::StringRef symName, mlir::Type type) {
+  if (mlir::isa<mlir::LLVM::LLVMPointerType>(type)) {
+    if (globalAS != programAS) {
+      auto llvmAddrOp = rewriter.create<mlir::LLVM::AddressOfOp>(
+          loc, getLlvmPtrType(rewriter.getContext(), globalAS), symName);
+      return rewriter.create<mlir::LLVM::AddrSpaceCastOp>(
+          loc, getLlvmPtrType(rewriter.getContext(), programAS), llvmAddrOp);
+    }
+    return rewriter.create<mlir::LLVM::AddressOfOp>(
+        loc, getLlvmPtrType(rewriter.getContext(), globalAS), symName);
+  }
+  return rewriter.create<mlir::LLVM::AddressOfOp>(loc, type, symName);
+}
+
+// Replaces an existing operation with an AddressOfOp or an AddrSpaceCastOp
+// depending on the existing address spaces of the type.
+mlir::Value replaceWithAddrOfOrASCast(mlir::ConversionPatternRewriter &rewriter,
+                                      mlir::Location loc,
+                                      std::uint64_t globalAS,
+                                      std::uint64_t programAS,
+                                      llvm::StringRef symName, mlir::Type type,
+                                      mlir::Operation *replaceOp) {
+  if (mlir::isa<mlir::LLVM::LLVMPointerType>(type)) {
+    if (globalAS != programAS) {
+      auto llvmAddrOp = rewriter.create<mlir::LLVM::AddressOfOp>(
+          loc, getLlvmPtrType(rewriter.getContext(), globalAS), symName);
+      return rewriter.replaceOpWithNewOp<mlir::LLVM::AddrSpaceCastOp>(
+          replaceOp, ::getLlvmPtrType(rewriter.getContext(), programAS),
+          llvmAddrOp);
+    }
+    return rewriter.replaceOpWithNewOp<mlir::LLVM::AddressOfOp>(
+        replaceOp, getLlvmPtrType(rewriter.getContext(), globalAS), symName);
+  }
+  return rewriter.replaceOpWithNewOp<mlir::LLVM::AddressOfOp>(replaceOp, type,
+                                                              symName);
----------------
skatrak wrote:

Nit: Feel free to disagree, but it seems like the logic here could be simplified a bit with something like this:
```suggestion
  bool isPointerType = mlir::isa<mlir::LLVM::LLVMPointerType>(type);
  if (isPointerType)
    type = getLlvmPtrType(rewriter.getContext(), globalAS)

  auto llvmAddrOp = rewriter.create<mlir::LLVM::AddressOfOp>(
      loc, type, symName);

  if (isPointerType && globalAS != programAS)
    return rewriter.replaceOpWithNewOp<mlir::LLVM::AddrSpaceCastOp>(
      replaceOp, getLlvmPtrType(rewriter.getContext(), programAS),
      llvmAddrOp);

  return rewriter.replaceOp(replaceOp, llvmAddrOp);
```

https://github.com/llvm/llvm-project/pull/119585