Fix broken _CCCL_BUILTIN_ASSUME macro (#3314)

* add compiler-specific path * fix device code path * add _CCC_ASSUME
NVIDIA · Jan 15, 2025 · f618ce4 · f618ce4
1 parent 43fb061
commit f618ce4
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/libcudacxx/include/cuda/std/__cccl/attributes.h b/libcudacxx/include/cuda/std/__cccl/attributes.h
@@ -134,4 +134,10 @@
 #  define _CCCL_RESTRICT __restrict__
 #endif // ^^^ !_CCCL_COMPILER(MSVC) ^^^
 
+#if _CCCL_HAS_CPP_ATTRIBUTE(assume)
+#  define _CCCL_ASSUME(...) [[assume(__VA_ARGS__)]]
+#else // ^^^ _CCCL_COMPILER(MSVC) ^^^ / vvv !_CCCL_COMPILER(MSVC) vvv
+#  define _CCCL_ASSUME(...) _CCCL_BUILTIN_ASSUME(__VA_ARGS__)
+#endif // ^^^ !_CCCL_COMPILER(MSVC) ^^^
+
 #endif // __CCCL_ATTRIBUTES_H
diff --git a/libcudacxx/include/cuda/std/__cccl/builtin.h b/libcudacxx/include/cuda/std/__cccl/builtin.h
@@ -101,8 +101,14 @@
 #  define _CCCL_BUILTIN_ADDRESSOF(...) __builtin_addressof(__VA_ARGS__)
 #endif // _CCCL_CHECK_BUILTIN(builtin_addressof)
 
-#if _CCCL_CHECK_BUILTIN(builtin_assume)
+#if _CCCL_CHECK_BUILTIN(builtin_assume) || _CCCL_COMPILER(CLANG) || _CCCL_COMPILER(NVHPC)
 #  define _CCCL_BUILTIN_ASSUME(...) __builtin_assume(__VA_ARGS__)
+#elif _CCCL_COMPILER(GCC, >=, 13)
+#  define _CCCL_BUILTIN_ASSUME(...) \
+    NV_IF_ELSE_TARGET(NV_IS_DEVICE, (__builtin_assume(__VA_ARGS__);), (__attribute__((__assume__(__VA_ARGS__)));))
+#elif _CCCL_COMPILER(MSVC)
+#  define _CCCL_BUILTIN_ASSUME(...) \
+    NV_IF_ELSE_TARGET(NV_IS_DEVICE, (__builtin_assume(__VA_ARGS__);), (__assume(__VA_ARGS__);))
 #endif // _CCCL_CHECK_BUILTIN(builtin_assume)
 
 // NVCC prior to 11.2 cannot handle __builtin_assume