From 1f4f410dcf0da671b3065bdb6408b4852c20b93f Mon Sep 17 00:00:00 2001
From: Sophia Wisdom <sophiawisdom1999@gmail.com>
Date: Tue, 15 Aug 2023 18:23:37 -0700
Subject: [PATCH] Update logging to give correct line/file/func numbers

fixes #376

Signed-off-by: Sophia Wisdom <sophiawisdom1999@gmail.com>
---
 .../include/transformer_engine/logging.h      | 30 +++++++++++--------
 1 file changed, 18 insertions(+), 12 deletions(-)

diff --git a/transformer_engine/common/include/transformer_engine/logging.h b/transformer_engine/common/include/transformer_engine/logging.h
index 9ac0bbbde2..5eebf8f7d6 100644
--- a/transformer_engine/common/include/transformer_engine/logging.h
+++ b/transformer_engine/common/include/transformer_engine/logging.h
@@ -20,6 +20,12 @@
                                  " in function " + __func__ + ": " + x);                           \
     } while (false)
 
+#define NVTE_LIBRARY_ERROR(x, file, line, func) \
+    do { \
+        throw std::runtime_error(std::string(file ":") + std::to_string(line) +            \
+                                 " in function " + func + ": " + x);                           \
+    } while (false)
+
 #define NVTE_CHECK(x, ...)                                                                         \
     do {                                                                                           \
         if (!(x)) {                                                                                \
@@ -29,19 +35,19 @@
 
 namespace {
 
-inline void check_cuda_(cudaError_t status) {
+inline void check_cuda_(cudaError_t status, const char *file, const int line, const char *func) {
     if ( status != cudaSuccess ) {
-        NVTE_ERROR("CUDA Error: " + std::string(cudaGetErrorString(status)));
+        NVTE_LIBRARY_ERROR("CUDA Error: " + std::string(cudaGetErrorString(status)), file, line, func);
     }
 }
 
-inline void check_cublas_(cublasStatus_t status) {
+inline void check_cublas_(cublasStatus_t status, const char *file, const int line, const char *func) {
     if ( status != CUBLAS_STATUS_SUCCESS ) {
-        NVTE_ERROR("CUBLAS Error: " + std::string(cublasGetStatusString(status)));
+        NVTE_LIBRARY_ERROR("CUBLAS Error: " + std::string(cublasGetStatusString(status)), file, line, func);
     }
 }
 
-inline void check_cudnn_(cudnnStatus_t status) {
+inline void check_cudnn_(cudnnStatus_t status, const char *file, const int line, const char *func) {
     if ( status != CUDNN_STATUS_SUCCESS ) {
         std::string message;
         message.reserve(1024);
@@ -51,24 +57,24 @@ inline void check_cudnn_(cudnnStatus_t status) {
                     "For more information, enable cuDNN error logging "
                     "by setting CUDNN_LOGERR_DBG=1 and "
                     "CUDNN_LOGDEST_DBG=stderr in the environment.");
-        NVTE_ERROR(message);
+        NVTE_LIBRARY_ERROR(message, file, line, func);
     }
 }
 
-inline void check_nvrtc_(nvrtcResult status) {
+inline void check_nvrtc_(nvrtcResult status, const char *file, const int line, const char *func) {
     if ( status != NVRTC_SUCCESS ) {
-        NVTE_ERROR("NVRTC Error: " + std::string(nvrtcGetErrorString(status)));
+        NVTE_LIBRARY_ERROR("NVRTC Error: " + std::string(nvrtcGetErrorString(status)), file, line, func);
     }
 }
 
 }  // namespace
 
-#define NVTE_CHECK_CUDA(ans) { check_cuda_(ans); }
+#define NVTE_CHECK_CUDA(ans) { check_cuda_(ans, __FILE__, __LINE__, __func__); }
 
-#define NVTE_CHECK_CUBLAS(ans) { check_cublas_(ans); }
+#define NVTE_CHECK_CUBLAS(ans) { check_cublas_(ans, __FILE__, __LINE__, __func__); }
 
-#define NVTE_CHECK_CUDNN(ans) { check_cudnn_(ans); }
+#define NVTE_CHECK_CUDNN(ans) { check_cudnn_(ans, __FILE__, __LINE__, __func__); }
 
-#define NVTE_CHECK_NVRTC(ans) { check_nvrtc_(ans); }
+#define NVTE_CHECK_NVRTC(ans) { check_nvrtc_(ans, __FILE__, __LINE__, __func__); }
 
 #endif  // TRANSFORMER_ENGINE_LOGGING_H_