Refactor logging macros

setup.py

@@ -467,6 +467,7 @@ def setup_pytorch_extension() -> setuptools.Extension:
     include_dirs = [
         root_path / "transformer_engine" / "common" / "include",
         root_path / "transformer_engine" / "pytorch" / "csrc",
+        root_path / "transformer_engine",
         root_path / "3rdparty" / "cudnn-frontend" / "include",
     ]
 
@@ -539,6 +540,7 @@ def setup_paddle_extension() -> setuptools.Extension:
     include_dirs = [
         root_path / "transformer_engine" / "common" / "include",
         root_path / "transformer_engine" / "paddle" / "csrc",
+        root_path / "transformer_engine",
     ]
 
     # Compiler flags

tests/cpp/operator/test_cast_transpose.cu

@@ -4,16 +4,17 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include <transformer_engine/transpose.h>
-#include <transformer_engine/logging.h>
-#include <gtest/gtest.h>
-#include <cuda_runtime.h>
-#include <cuda_bf16.h>
-#include <memory>
-#include <iostream>
+#include <cstring>
 #include <iomanip>
+#include <iostream>
+#include <memory>
 #include <random>
-#include <cstring>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <gtest/gtest.h>
+
+#include <transformer_engine/transpose.h>
 #include "../test_common.h"
 
 using namespace transformer_engine;

tests/cpp/operator/test_cast_transpose_dbias.cu

@@ -4,17 +4,18 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include <transformer_engine/transpose.h>
-#include <transformer_engine/logging.h>
-#include <gtest/gtest.h>
-#include <cuda_runtime.h>
-#include <cuda_bf16.h>
+#include <cmath>
+#include <cstring>
 #include <memory>
-#include <iostream>
 #include <iomanip>
+#include <iostream>
 #include <random>
-#include <cstring>
-#include <cmath>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <gtest/gtest.h>
+
+#include <transformer_engine/transpose.h>
 #include "../test_common.h"
 
 using namespace transformer_engine;

tests/cpp/operator/test_cast_transpose_dbias_dgelu.cu

@@ -4,17 +4,18 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include <transformer_engine/transpose.h>
-#include <transformer_engine/logging.h>
-#include <gtest/gtest.h>
-#include <cuda_runtime.h>
-#include <cuda_bf16.h>
+#include <cmath>
+#include <cstring>
 #include <memory>
-#include <iostream>
 #include <iomanip>
+#include <iostream>
 #include <random>
-#include <cstring>
-#include <cmath>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <gtest/gtest.h>
+
+#include <transformer_engine/transpose.h>
 #include "../test_common.h"
 
 using namespace transformer_engine;

tests/cpp/operator/test_cast_transpose_dgeglu.cu

@@ -4,17 +4,18 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-#include <gtest/gtest.h>
-#include <transformer_engine/logging.h>
-#include <transformer_engine/transpose.h>
 #include <cmath>
 #include <cstring>
 #include <iomanip>
 #include <iostream>
 #include <memory>
 #include <random>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <gtest/gtest.h>
+
+#include <transformer_engine/transpose.h>
 #include "../test_common.h"
 
 using namespace transformer_engine;

tests/cpp/operator/test_dgeglu.cu

@@ -4,18 +4,19 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-#include <gtest/gtest.h>
-#include <transformer_engine/activation.h>
-#include <transformer_engine/logging.h>
 #include <cmath>
 #include <cstring>
 #include <iomanip>
 #include <iostream>
 #include <memory>
 #include <random>
 #include <type_traits>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <gtest/gtest.h>
+
+#include <transformer_engine/activation.h>
 #include "../test_common.h"
 
 using namespace transformer_engine;

tests/cpp/operator/test_geglu.cu

@@ -4,18 +4,19 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-#include <gtest/gtest.h>
-#include <transformer_engine/activation.h>
-#include <transformer_engine/logging.h>
 #include <cmath>
 #include <cstring>
 #include <iomanip>
 #include <iostream>
 #include <memory>
 #include <random>
 #include <type_traits>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <gtest/gtest.h>
+
+#include <transformer_engine/activation.h>
 #include "../test_common.h"
 
 using namespace transformer_engine;

tests/cpp/operator/test_gelu.cu

@@ -4,18 +4,19 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include <transformer_engine/activation.h>
-#include <transformer_engine/logging.h>
-#include <gtest/gtest.h>
-#include <cuda_runtime.h>
-#include <cuda_bf16.h>
 #include <cmath>
+#include <cstring>
 #include <memory>
-#include <iostream>
 #include <iomanip>
+#include <iostream>
 #include <random>
-#include <cstring>
 #include <type_traits>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <gtest/gtest.h>
+
+#include <transformer_engine/activation.h>
 #include "../test_common.h"
 
 using namespace transformer_engine;

tests/cpp/operator/test_layernorm.cu

@@ -4,17 +4,19 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include <transformer_engine/layer_norm.h>
-#include <transformer_engine/transformer_engine.h>
-#include <gtest/gtest.h>
-#include <cuda_runtime.h>
-#include <cuda_bf16.h>
+#include <cmath>
+#include <cstring>
 #include <memory>
-#include <iostream>
 #include <iomanip>
+#include <iostream>
 #include <random>
-#include <cstring>
-#include <cmath>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <gtest/gtest.h>
+
+#include <transformer_engine/layer_norm.h>
+#include <transformer_engine/transformer_engine.h>
 #include "../test_common.h"
 
 using namespace transformer_engine;

tests/cpp/operator/test_multi_cast_transpose.cu

@@ -4,17 +4,18 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include <transformer_engine/transpose.h>
-#include <transformer_engine/logging.h>
-#include <gtest/gtest.h>
-#include <cuda_runtime.h>
-#include <cuda_bf16.h>
 #include <cstring>
-#include <iostream>
 #include <iomanip>
+#include <iostream>
 #include <memory>
 #include <random>
 #include <vector>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <gtest/gtest.h>
+
+#include <transformer_engine/transpose.h>
 #include "../test_common.h"
 
 using namespace transformer_engine;

tests/cpp/operator/test_qdq.cu

@@ -4,19 +4,19 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include "gtest/gtest.h"
-#include <transformer_engine/cast.h>
-#include <transformer_engine/logging.h>
-#include <gtest/gtest.h>
-#include <cuda_runtime.h>
-#include <cuda_bf16.h>
-#include <memory>
-#include <iostream>
+#include <cstring>
 #include <iomanip>
+#include <iostream>
+#include <memory>
 #include <random>
-#include <cstring>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <gtest/gtest.h>
+
+#include <transformer_engine/cast.h>
+#include <transformer_engine/transformer_engine.h>
 #include "../test_common.h"
-#include "transformer_engine/transformer_engine.h"
 
 using namespace transformer_engine;
 

tests/cpp/operator/test_rmsnorm.cu

@@ -4,17 +4,19 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-#include <gtest/gtest.h>
-#include <transformer_engine/rmsnorm.h>
-#include <transformer_engine/transformer_engine.h>
 #include <cmath>
 #include <cstring>
 #include <iomanip>
 #include <iostream>
 #include <memory>
 #include <random>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <gtest/gtest.h>
+
+#include <transformer_engine/rmsnorm.h>
+#include <transformer_engine/transformer_engine.h>
 #include "../test_common.h"
 
 using namespace transformer_engine;

tests/cpp/operator/test_transpose.cu

@@ -4,16 +4,17 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include <transformer_engine/transpose.h>
-#include <transformer_engine/logging.h>
-#include <gtest/gtest.h>
-#include <cuda_runtime.h>
-#include <cuda_bf16.h>
-#include <memory>
-#include <iostream>
+#include <cstring>
 #include <iomanip>
+#include <iostream>
+#include <memory>
 #include <random>
-#include <cstring>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <gtest/gtest.h>
+
+#include <transformer_engine/transpose.h>
 #include "../test_common.h"
 
 using namespace transformer_engine;

tests/cpp/test_common.cu

@@ -6,13 +6,16 @@
 
 
 #include "test_common.h"
-#include "transformer_engine/logging.h"
-#include "transformer_engine/transformer_engine.h"
-#include <gtest/gtest.h>
+
 #include <algorithm>
 #include <memory>
 #include <random>
 
+#include <gtest/gtest.h>
+
+#include <transformer_engine/transformer_engine.h>
+#include "util/logging.h"
+
 namespace test {
 
 std::vector<DType> all_fp_types = {DType::kFloat32,

tests/cpp/test_common.h

@@ -6,15 +6,17 @@
 
 #pragma once
 
+#include <iostream>
 #include <memory>
-#include <transformer_engine/transformer_engine.h>
-#include <transformer_engine/logging.h>
-#include <cuda_fp16.h>
+#include <vector>
+
 #include <cuda_bf16.h>
+#include <cuda_fp16.h>
 #include <cuda_fp8.h>
 #include <cuda_runtime_api.h>
-#include <vector>
-#include <iostream>
+
+#include <transformer_engine/transformer_engine.h>
+#include "util/logging.h"
 
 namespace test {
 using namespace transformer_engine;
@@ -252,4 +254,3 @@ bool isFp8Type(DType type);
         default: \
             NVTE_ERROR("Invalid type."); \
     }
-