oom error

Author: drisspg

CMakeLists.txt

@@ -42,18 +42,34 @@ option (SHOW_PTXAS_INFO "Show ptxas info" OFF)
 if(SHOW_PTXAS_INFO)
     set(EXTRA_CUDA_FLAGS "${EXTRA_CUDA_FLAGS} -Xptxas -v")
 endif()
+
+set(SIMPLE_CUDA_ARCH "121" CACHE STRING "CUDA architecture (e.g. 121, 100a)")
+set(SIMPLE_CUDA_SKIP_EXAMPLES "" CACHE STRING "Semicolon-separated example target names to skip")
+set(SIMPLE_CUDA_AUTO_SKIP_UNSUPPORTED_EXAMPLES ON CACHE BOOL "Auto-skip examples unsupported on selected architecture")
+
+set(EFFECTIVE_SKIP_EXAMPLES ${SIMPLE_CUDA_SKIP_EXAMPLES})
+if(SIMPLE_CUDA_AUTO_SKIP_UNSUPPORTED_EXAMPLES AND SIMPLE_CUDA_ARCH MATCHES "(^|;)121($|;)")
+    list(APPEND EFFECTIVE_SKIP_EXAMPLES stochastic test_nvfp4_rounding)
+endif()
+list(REMOVE_DUPLICATES EFFECTIVE_SKIP_EXAMPLES)
 # Set the output directory for the binaries
 set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
 
 # Build the examples
 file(GLOB_RECURSE EXAMPLE_SOURCES examples/*.cu)
 foreach(EXAMPLE_SOURCE ${EXAMPLE_SOURCES})
     get_filename_component(EXAMPLE_NAME ${EXAMPLE_SOURCE} NAME_WE)
+    list(FIND EFFECTIVE_SKIP_EXAMPLES ${EXAMPLE_NAME} SKIP_EXAMPLE_INDEX)
+    if(NOT SKIP_EXAMPLE_INDEX EQUAL -1)
+        message(STATUS "Skipping example ${EXAMPLE_NAME}")
+        continue()
+    endif()
+
     add_executable(${EXAMPLE_NAME} ${EXAMPLE_SOURCE})
 
     # CUDA properties provided by CMAKE
     set_target_properties(${EXAMPLE_NAME} PROPERTIES CUDA_SEPARABLE_COMPILATION ON)
-    set_target_properties(${EXAMPLE_NAME} PROPERTIES CUDA_ARCHITECTURES 100a)
+    set_target_properties(${EXAMPLE_NAME} PROPERTIES CUDA_ARCHITECTURES ${SIMPLE_CUDA_ARCH})
 
     # Convert the flags string into a list of flags
     separate_arguments(EXTRA_CUDA_FLAGS_LIST UNIX_COMMAND "${EXTRA_CUDA_FLAGS}")

examples/misc/mem_bw.cu

@@ -37,7 +37,8 @@ int main() {
   std::fill_n(data, n, 1); // initialize data
 
   // Make sure data is on the device before timing
-  cudaMemPrefetchAsync(data, n * sizeof(float), 0);
+  cudaMemLocation host_location{cudaMemLocationTypeHost, 0};
+  cudaMemPrefetchAsync(data, n * sizeof(float), host_location, 0, 0);
   cudaDeviceSynchronize();
 
   direct_copy_optimized<<<nBlocks_manual, blockSize>>>(reinterpret_cast<float4*>(output), reinterpret_cast<float4*>(data), n);

examples/misc/oom_400gb.cu

@@ -0,0 +1,33 @@
+#include <cuda_runtime.h>
+#include <fmt/core.h>
+#include <cstdlib>
+#include <cstdint>
+
+int main() {
+  constexpr std::size_t gib = 1024ULL * 1024ULL * 1024ULL;
+  constexpr std::size_t requested_bytes = 400ULL * gib;
+
+  fmt::print("Requesting {} GiB ({} bytes)\n", requested_bytes / gib, requested_bytes);
+
+  void* host_ptr = std::malloc(requested_bytes);
+  if (host_ptr == nullptr) {
+    fmt::print("malloc: OOM (nullptr)\n");
+  } else {
+    fmt::print("malloc: success (can still be overcommit)\n");
+    std::free(host_ptr);
+  }
+
+  void* pinned_ptr = nullptr;
+  cudaError_t pinned_alloc_result = cudaMallocHost(&pinned_ptr, requested_bytes);
+  if (pinned_alloc_result == cudaSuccess) {
+    fmt::print("cudaMallocHost: success\n");
+    cudaFreeHost(pinned_ptr);
+  } else {
+    fmt::print(
+        "cudaMallocHost: {} ({})\n",
+        cudaGetErrorName(pinned_alloc_result),
+        cudaGetErrorString(pinned_alloc_result));
+  }
+
+  return 0;
+}

examples/test_nvfp4_rounding.cu

@@ -9,6 +9,7 @@
 #include <vector>
 #include <cmath>
 #include <algorithm>
+#include <cstdint>
 
 // FP4 E2M1 lookup table
 const float fp4_e2m1_lut[16] = {