[OpenMP][Offload] Add support for lambdas with debug conditions

[adurang]

This PR adds a new set of debug macros that allow a certain code to be only executed when certain debug conditions are met. This is useful to guard things that are not strictly messages but compute and store things that are related to those messages.

Strictly speaking the existing ODBG_OS could be used as well but that requires a stream object to be created which is unnecessary in some cases.

Example of how it works:

ODBG_IF("Counters", [&](uint32_t Level) {
    someCounter++;
    if (Level == 2) moreDetailedCounter += f();
});
ODBG("Counters") << "Counter" =  someCounter 
                 << ODBG_IF(2) << "DetailedCounter" << moreDetailedCounter;

[llvmbot]

@llvm/pr-subscribers-offload

Author: Alex Duran (adurang)

Changes

This PR adds a new set of debug macros that allow a certain code to be only executed when certain debug conditions are met. This is useful to guard things that are not strictly messages but compute and store things that are related to those messages.

Strictly speaking the existing ODBG_OS could be used as well but that requires a stream object to be created which is unnecessary in some cases.

Example of how it works:

ODBG_IF("Counters", [&](uint32_t Level) {
someCounter++;
if (Level == 2) moreDetailedCounter += f();
});
ODBG("Counters") << "Counter" = someCounter

---

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

1 Files Affected:

• (modified) offload/include/Shared/Debug.h (+46-9)

diff --git a/offload/include/Shared/Debug.h b/offload/include/Shared/Debug.h
index d5cf719f1ebf3..2ff0eb80788b8 100644
--- a/offload/include/Shared/Debug.h
+++ b/offload/include/Shared/Debug.h
@@ -432,15 +432,22 @@ static inline raw_ostream &operator<<(raw_ostream &Os,
 
 // helper templates to support lambdas with different number of arguments
 template <typename LambdaTy> struct LambdaHelper {
-  template <typename T, typename = std::void_t<>>
-  struct has_two_args : std::false_type {};
-  template <typename T>
-  struct has_two_args<T,
-                      std::void_t<decltype(std::declval<T>().operator()(1, 2))>>
-      : std::true_type {};
+  template <typename FuncTy, typename RetTy, typename... Args>
+  static constexpr size_t CountArgs(RetTy (FuncTy::*)(Args...)) {
+    return sizeof...(Args);
+  }
+  template <typename FuncTy, typename RetTy, typename... Args>
+  static constexpr size_t CountArgs(RetTy (FuncTy::*)(Args...) const) {
+    return sizeof...(Args);
+  }
 
+  static constexpr size_t NArgs = CountArgs(&LambdaTy::operator());
+}
+
+template <typename LambdaTy>
+struct LambdaOs : public LambdaHelper<LambdaTy> {
   static void dispatch(LambdaTy func, llvm::raw_ostream &Os, uint32_t Level) {
-    if constexpr (has_two_args<LambdaTy>::value)
+    if constexpr (NArgs == 2)
       func(Os, Level);
     else
       func(Os);
@@ -457,8 +464,8 @@ template <typename LambdaTy> struct LambdaHelper {
           RealLevel, /*ShouldPrefixNextString=*/true,                          \
           /*ShouldEmitNewLineOnDestruction=*/true};                            \
       auto F = Callback;                                                       \
-      ::llvm::offload::debug::LambdaHelper<decltype(F)>::dispatch(F, OS,       \
-                                                                  RealLevel);  \
+      ::llvm::offload::debug::LambdaOs<decltype(F)>::dispatch(F, OS,           \
+                                                              RealLevel);      \
     }                                                                          \
   }
 
@@ -476,6 +483,33 @@ template <typename LambdaTy> struct LambdaHelper {
 #define ODBG_OS(...)                                                           \
   ODBG_OS_SELECT(__VA_ARGS__ __VA_OPT__(, ) 3, 2, 1)(__VA_ARGS__)
 
+// helper templates to support lambdas with different number of arguments
+template <typename LambdaTy> struct LambdaIf : public LambdaHelper<LambdaTy> {
+  static void dispatch(LambdaTy func, uint32_t Level) {
+    if constexpr (NArgs == 1)
+      func(Level);
+    else
+      func();
+  }
+};
+
+#define ODBG_IF_BASE(Type, Level, Callback)                                    \
+  if (::llvm::offload::debug::isDebugEnabled()) {                              \
+    uint32_t RealLevel = (Level);                                              \
+    if (::llvm::offload::debug::shouldPrintDebug(GETNAME(TARGET_NAME), (Type), \
+                                                 RealLevel)) {                 \
+      auto F = Callback;                                                       \
+      ::llvm::offload::debug::LambdaIf<decltype(F)>::dispatch(F, RealLevel);   \
+    }                                                                          \
+  }
+
+#define ODBG_IF_3(Type, Level, Callback) ODBG_IF_BASE(Type, Level, Callback)
+#define ODBG_IF_2(Type, Callback) ODBG_IF_3(Type, 1, Callback)
+#define ODBG_IF_1(Callback) ODBG_IF_2("default", Callback)
+#define ODBG_IF_SELECT(Type, Level, Callback, NArgs, ...) ODBG_IF_##NArgs
+#define ODBG_IF(...)                                                           \
+  ODBG_IF_SELECT(__VA_ARGS__ __VA_OPT__(, ) 3, 2, 1)(__VA_ARGS__)
+
 #else
 
 inline bool isDebugEnabled() { return false; }
@@ -496,6 +530,9 @@ inline bool isDebugEnabled() { return false; }
 #define ODBG_OS_STREAM(Stream, Type, Level, Callback)
 #define ODBG_OS(...)
 
+#define ODBG_IF_BASE(Type, Level, Callback)
+#define ODBG_IF(...)
+
 #endif
 
 } // namespace llvm::offload::debug

[adurang]

@mjklemm I decided to restore the helper code we removed in the previous PR as I think it makes more sense now to use that instead of having to create different "has_xx_args" helpers for the different cases.