From 4ed302952b423a76f964b826b45d2d504c5584b8 Mon Sep 17 00:00:00 2001
From: Allard Hendriksen <ahendriksen@nvidia.com>
Date: Tue, 4 Apr 2023 16:10:47 +0200
Subject: [PATCH] Move RMM_LOGGING_ASSERT into separate header

Avoiding the inclusion of rmm/logger.hpp can reduce compile times
significantly.
---
 include/rmm/cuda_stream.hpp                   |  2 +
 include/rmm/detail/error.hpp                  | 23 ---------
 include/rmm/detail/logging_assert.hpp         | 47 +++++++++++++++++++
 .../rmm/mr/device/arena_memory_resource.hpp   |  1 +
 include/rmm/mr/device/detail/arena.hpp        |  1 +
 .../mr/device/fixed_size_memory_resource.hpp  |  1 +
 .../rmm/mr/device/pool_memory_resource.hpp    |  1 +
 .../mr/device/tracking_resource_adaptor.hpp   |  1 +
 tests/mr/device/tracking_mr_tests.cpp         |  1 +
 9 files changed, 55 insertions(+), 23 deletions(-)
 create mode 100644 include/rmm/detail/logging_assert.hpp

diff --git a/include/rmm/cuda_stream.hpp b/include/rmm/cuda_stream.hpp
index 185cd049e..f80ba8969 100644
--- a/include/rmm/cuda_stream.hpp
+++ b/include/rmm/cuda_stream.hpp
@@ -16,8 +16,10 @@
 
 #pragma once
 
+#include <functional>
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/detail/error.hpp>
+#include <rmm/detail/logging_assert.hpp>
 
 #include <cuda_runtime_api.h>
 
diff --git a/include/rmm/detail/error.hpp b/include/rmm/detail/error.hpp
index 1c669822c..5d8b340ab 100644
--- a/include/rmm/detail/error.hpp
+++ b/include/rmm/detail/error.hpp
@@ -16,8 +16,6 @@
 
 #pragma once
 
-#include <rmm/logger.hpp>
-
 #include <cuda_runtime_api.h>
 
 #include <cassert>
@@ -245,24 +243,3 @@ class out_of_range : public std::out_of_range {
     assert(status__ == cudaSuccess);                                            \
   } while (0)
 #endif
-
-/**
- * @brief Assertion that logs a CRITICAL log message on failure.
- */
-#ifdef NDEBUG
-#define RMM_LOGGING_ASSERT(_expr) (void)0
-#elif SPDLOG_ACTIVE_LEVEL < SPDLOG_LEVEL_OFF
-#define RMM_LOGGING_ASSERT(_expr)                                                                 \
-  do {                                                                                            \
-    bool const success = (_expr);                                                                 \
-    if (!success) {                                                                               \
-      RMM_LOG_CRITICAL(                                                                           \
-        "[" __FILE__ ":" RMM_STRINGIFY(__LINE__) "] Assertion " RMM_STRINGIFY(_expr) " failed."); \
-      rmm::logger().flush();                                                                      \
-      /* NOLINTNEXTLINE(cppcoreguidelines-pro-bounds-array-to-pointer-decay) */                   \
-      assert(success);                                                                            \
-    }                                                                                             \
-  } while (0)
-#else
-#define RMM_LOGGING_ASSERT(_expr) assert((_expr));
-#endif
diff --git a/include/rmm/detail/logging_assert.hpp b/include/rmm/detail/logging_assert.hpp
new file mode 100644
index 000000000..321bddab8
--- /dev/null
+++ b/include/rmm/detail/logging_assert.hpp
@@ -0,0 +1,47 @@
+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+// Only include <rmm/logger.hpp> if needed in RMM_LOGGING_ASSERT below. The
+// logger can be extremely expensive to compile, so we want to avoid including
+// it.
+#if !defined(NDEBUG)
+#include <cassert>
+#include <rmm/detail/error.hpp>
+#include <rmm/logger.hpp>
+#endif
+
+/**
+ * @brief Assertion that logs a CRITICAL log message on failure.
+ */
+#ifdef NDEBUG
+#define RMM_LOGGING_ASSERT(_expr) (void)0
+#elif SPDLOG_ACTIVE_LEVEL < SPDLOG_LEVEL_OFF
+#define RMM_LOGGING_ASSERT(_expr)                                                                 \
+  do {                                                                                            \
+    bool const success = (_expr);                                                                 \
+    if (!success) {                                                                               \
+      RMM_LOG_CRITICAL(                                                                           \
+        "[" __FILE__ ":" RMM_STRINGIFY(__LINE__) "] Assertion " RMM_STRINGIFY(_expr) " failed."); \
+      rmm::logger().flush();                                                                      \
+      /* NOLINTNEXTLINE(cppcoreguidelines-pro-bounds-array-to-pointer-decay) */                   \
+      assert(success);                                                                            \
+    }                                                                                             \
+  } while (0)
+#else
+#define RMM_LOGGING_ASSERT(_expr) assert((_expr));
+#endif
diff --git a/include/rmm/mr/device/arena_memory_resource.hpp b/include/rmm/mr/device/arena_memory_resource.hpp
index b007d8f54..123afb9ae 100644
--- a/include/rmm/mr/device/arena_memory_resource.hpp
+++ b/include/rmm/mr/device/arena_memory_resource.hpp
@@ -16,6 +16,7 @@
 #pragma once
 
 #include <rmm/detail/error.hpp>
+#include <rmm/detail/logging_assert.hpp>
 #include <rmm/logger.hpp>
 #include <rmm/mr/device/detail/arena.hpp>
 #include <rmm/mr/device/device_memory_resource.hpp>
diff --git a/include/rmm/mr/device/detail/arena.hpp b/include/rmm/mr/device/detail/arena.hpp
index 493fb6cb1..de5cbb4a5 100644
--- a/include/rmm/mr/device/detail/arena.hpp
+++ b/include/rmm/mr/device/detail/arena.hpp
@@ -20,6 +20,7 @@
 #include <rmm/detail/aligned.hpp>
 #include <rmm/detail/cuda_util.hpp>
 #include <rmm/detail/error.hpp>
+#include <rmm/detail/logging_assert.hpp>
 #include <rmm/logger.hpp>
 
 #include <cuda_runtime_api.h>
diff --git a/include/rmm/mr/device/fixed_size_memory_resource.hpp b/include/rmm/mr/device/fixed_size_memory_resource.hpp
index 9a1923c2d..fda4625ac 100644
--- a/include/rmm/mr/device/fixed_size_memory_resource.hpp
+++ b/include/rmm/mr/device/fixed_size_memory_resource.hpp
@@ -18,6 +18,7 @@
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/detail/aligned.hpp>
 #include <rmm/detail/error.hpp>
+#include <rmm/detail/logging_assert.hpp>
 #include <rmm/mr/device/detail/fixed_size_free_list.hpp>
 #include <rmm/mr/device/detail/stream_ordered_memory_resource.hpp>
 
diff --git a/include/rmm/mr/device/pool_memory_resource.hpp b/include/rmm/mr/device/pool_memory_resource.hpp
index 297b3f864..beb7996e8 100644
--- a/include/rmm/mr/device/pool_memory_resource.hpp
+++ b/include/rmm/mr/device/pool_memory_resource.hpp
@@ -19,6 +19,7 @@
 #include <rmm/detail/aligned.hpp>
 #include <rmm/detail/cuda_util.hpp>
 #include <rmm/detail/error.hpp>
+#include <rmm/detail/logging_assert.hpp>
 #include <rmm/logger.hpp>
 #include <rmm/mr/device/detail/coalescing_free_list.hpp>
 #include <rmm/mr/device/detail/stream_ordered_memory_resource.hpp>
diff --git a/include/rmm/mr/device/tracking_resource_adaptor.hpp b/include/rmm/mr/device/tracking_resource_adaptor.hpp
index f1996ba01..8ce2c439e 100644
--- a/include/rmm/mr/device/tracking_resource_adaptor.hpp
+++ b/include/rmm/mr/device/tracking_resource_adaptor.hpp
@@ -17,6 +17,7 @@
 
 #include <rmm/detail/error.hpp>
 #include <rmm/detail/stack_trace.hpp>
+#include <rmm/logger.hpp>
 #include <rmm/mr/device/device_memory_resource.hpp>
 
 #include <fmt/core.h>
diff --git a/tests/mr/device/tracking_mr_tests.cpp b/tests/mr/device/tracking_mr_tests.cpp
index e804cf0ea..2af720d3f 100644
--- a/tests/mr/device/tracking_mr_tests.cpp
+++ b/tests/mr/device/tracking_mr_tests.cpp
@@ -19,6 +19,7 @@
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/detail/error.hpp>
 #include <rmm/device_buffer.hpp>
+#include <rmm/logger.hpp>
 #include <rmm/mr/device/tracking_resource_adaptor.hpp>
 
 #include <gtest/gtest.h>