[llvm] [Offload]: Skip copying of unused kernel-mapped data (PR #124723)

[via llvm-commits]

https://github.com/pradt2 updated https://github.com/llvm/llvm-project/pull/124723

>From 0e8c62a298759464c02fb33d11f74567b7492721 Mon Sep 17 00:00:00 2001
From: pradt2 <12902844+pradt2 at users.noreply.github.com>
Date: Tue, 28 Jan 2025 01:00:58 -0800
Subject: [PATCH] [Offload]: Skip copying of unused kernel-mapped data

---
 offload/src/omptarget.cpp | 38 ++++++++++++++++++++++++++++++++++++--
 1 file changed, 36 insertions(+), 2 deletions(-)

diff --git a/offload/src/omptarget.cpp b/offload/src/omptarget.cpp
index 5b25d955dd320d..729669caa07bc0 100644
--- a/offload/src/omptarget.cpp
+++ b/offload/src/omptarget.cpp
@@ -1197,6 +1197,35 @@ class PrivateArgumentManagerTy {
   }
 };
 
+static std::unique_ptr<int64_t[]> maskIgnorableMappings(int64_t DeviceId, int32_t ArgNum, int64_t *ArgTypes,
+                                                        int64_t *ArgSizes, map_var_info_t *ArgNames) {
+    std::unique_ptr<int64_t[]> ArgTypesOverride = std::make_unique<int64_t[]>(ArgNum);
+
+    for (int32_t I = 0; I < ArgNum; ++I) {
+        bool IsTargetParam = ArgTypes[I] & OMP_TGT_MAPTYPE_TARGET_PARAM;
+
+        bool IsMapTo = ArgTypes[I] & OMP_TGT_MAPTYPE_TO;
+        if (IsTargetParam || !IsMapTo) {
+            ArgTypesOverride[I] = ArgTypes[I];
+            continue;
+        }
+
+        bool IsMapFrom = ArgTypes[I] & OMP_TGT_MAPTYPE_FROM;
+        const char *Type = IsMapFrom ? "tofrom" : "to";
+
+        // Optimisation: A 'to' or 'tofrom' mapping is not
+        // used by the kernel. Change its type such that
+        // no new mapping is created, but any existing
+        // mapping has its counter decremented.
+        INFO(OMP_INFOTYPE_ALL, DeviceId, "%s(%s)[%" PRId64 "] %s\n", Type,
+             getNameFromMapping(ArgNames[I]).c_str(), ArgSizes[I], "is not used and will not be copied");
+
+        ArgTypesOverride[I] = ArgTypes[I] & ~(OMP_TGT_MAPTYPE_TO | OMP_TGT_MAPTYPE_FROM);
+    }
+
+    return ArgTypesOverride;
+}
+
 /// Process data before launching the kernel, including calling targetDataBegin
 /// to map and transfer data to target device, transferring (first-)private
 /// variables.
@@ -1417,11 +1446,16 @@ int target(ident_t *Loc, DeviceTy &Device, void *HostPtr,
 
   int NumClangLaunchArgs = KernelArgs.NumArgs;
   int Ret = OFFLOAD_SUCCESS;
+
+    std::unique_ptr<int64_t[]> ArgTypesOverride =
+            maskIgnorableMappings(DeviceId, NumClangLaunchArgs, KernelArgs.ArgTypes,
+                                  KernelArgs.ArgSizes, KernelArgs.ArgNames);
+
   if (NumClangLaunchArgs) {
     // Process data, such as data mapping, before launching the kernel
     Ret = processDataBefore(Loc, DeviceId, HostPtr, NumClangLaunchArgs,
                             KernelArgs.ArgBasePtrs, KernelArgs.ArgPtrs,
-                            KernelArgs.ArgSizes, KernelArgs.ArgTypes,
+                            KernelArgs.ArgSizes, ArgTypesOverride.get(),
                             KernelArgs.ArgNames, KernelArgs.ArgMappers, TgtArgs,
                             TgtOffsets, PrivateArgumentManager, AsyncInfo);
     if (Ret != OFFLOAD_SUCCESS) {
@@ -1473,7 +1507,7 @@ int target(ident_t *Loc, DeviceTy &Device, void *HostPtr,
     // variables
     Ret = processDataAfter(Loc, DeviceId, HostPtr, NumClangLaunchArgs,
                            KernelArgs.ArgBasePtrs, KernelArgs.ArgPtrs,
-                           KernelArgs.ArgSizes, KernelArgs.ArgTypes,
+                           KernelArgs.ArgSizes, ArgTypesOverride.get(),
                            KernelArgs.ArgNames, KernelArgs.ArgMappers,
                            PrivateArgumentManager, AsyncInfo);
     if (Ret != OFFLOAD_SUCCESS) {