WIP

Author: wujingyue

CMakeLists.txt

@@ -808,6 +808,7 @@ if(BUILD_NVFUSER_BENCHMARK)
   target_link_libraries(nvfuser_bench PRIVATE
     benchmark::benchmark
     codegen_internal
+    GTest::gtest
   )
   add_dependencies(nvfuser_bench flatc build_flatbuffer_config)
 
@@ -845,6 +846,7 @@ if(BUILD_NVFUSER_BENCHMARK)
     target_link_libraries(nvfuser_multidevice_bench PRIVATE
       benchmark::benchmark
       codegen_internal
+      GTest::gtest
     )
     add_dependencies(nvfuser_multidevice_bench flatc build_flatbuffer_config)
 

tests/cpp/utils.cpp

@@ -46,10 +46,13 @@ void NVFuserTest::SetUp() {
   EnableOptionsGuard::getCurOptions().set(EnableOption::IdModelExtraValidation);
 
   resetPeakMemoryStats(0);
-  NVF_ERROR_EQ(maxMemoryAllocated(0), 0);
+  ASSERT_EQ(maxMemoryAllocated(0), 0);
 }
 
 void NVFuserTest::TearDown() {
+  at::cuda::clearCublasWorkspaces();
+
+  ASSERT_EQ(memoryAllocated(0), 0);
   if (::testing::Test::HasFailure()) {
     auto test_info = ::testing::UnitTest::GetInstance()->current_test_info();
     std::cerr << "To reproduce: NVFUSER_TEST_RANDOM_SEED=" << getCRandomSeed()
@@ -930,4 +933,17 @@ int64_t maxMemoryAllocated(const c10::DeviceIndex device) {
       .peak;
 }
 
+int64_t memoryAllocated(const c10::DeviceIndex device) {
+  c10::cuda::CUDACachingAllocator::CUDAAllocator* allocator =
+      c10::cuda::CUDACachingAllocator::get();
+  NVF_CHECK(allocator != nullptr);
+
+  c10::CachingDeviceAllocator::DeviceStats device_stats =
+      allocator->getDeviceStats(device);
+
+  return device_stats.allocated_bytes
+      .at(static_cast<uint64_t>(c10::CachingAllocator::StatType::AGGREGATE))
+      .current;
+}
+
 } // namespace nvfuser

tests/cpp/utils.h

@@ -847,4 +847,7 @@ void resetPeakMemoryStats(c10::DeviceIndex device);
 // C++ implementation of torch.cuda.max_memory_allocated
 int64_t maxMemoryAllocated(const c10::DeviceIndex device);
 
+// C++ implementation of torch.cuda.memory_allocated
+int64_t memoryAllocated(const c10::DeviceIndex device);
+
 } // namespace nvfuser

tests/cpp/validator.cpp

@@ -6,8 +6,8 @@
  */
 // clang-format on
 #include <runtime/fusion_kernel_runtime.h>
+#include <tests/cpp/utils.h>
 #include <tests/cpp/validator.h>
-
 #include <validator_utils.h>
 
 namespace nvfuser {