Sync ptx_dot_variants.h with libcuda-ptx (NVIDIA#3564)

davebayer · Jan 29, 2025 · 59b05c5 · 59b05c5
1 parent dc5f31f
commit 59b05c5
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 3 deletions.
diff --git a/libcudacxx/include/cuda/__ptx/ptx_dot_variants.h b/libcudacxx/include/cuda/__ptx/ptx_dot_variants.h
@@ -9,6 +9,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+// WARNING: The source of truth for this file is libcuda-ptx. Do not modify without syncing with libcuda-ptx.
+
 #ifndef _CUDA_PTX_DOT_VARIANTS_H_
 #define _CUDA_PTX_DOT_VARIANTS_H_
 
@@ -111,6 +113,23 @@ enum class dot_op
   exch
 };
 
+enum class dot_cta_group
+{
+  cta_group_1,
+  cta_group_2
+};
+
+enum class dot_kind
+{
+  f16,
+  f8f6f4,
+  i8,
+  mxf4,
+  mxf4nvf4,
+  mxf8f6f4,
+  tf32
+};
+
 template <dot_sem __sem>
 using sem_t         = _CUDA_VSTD::integral_constant<dot_sem, __sem>;
 using sem_acq_rel_t = sem_t<dot_sem::acq_rel>;
@@ -173,6 +192,35 @@ static constexpr op_xor_op_t op_xor_op{};
 static constexpr op_cas_t op_cas{};
 static constexpr op_exch_t op_exch{};
 
+template <dot_cta_group __cta_group>
+using cta_group_t   = _CUDA_VSTD::integral_constant<dot_cta_group, __cta_group>;
+using cta_group_1_t = cta_group_t<dot_cta_group::cta_group_1>;
+using cta_group_2_t = cta_group_t<dot_cta_group::cta_group_1>;
+
+static constexpr cta_group_1_t cta_group_1{};
+static constexpr cta_group_2_t cta_group_2{};
+
+template <dot_kind __kind>
+using kind_t          = _CUDA_VSTD::integral_constant<dot_kind, __kind>;
+using kind_f16_t      = kind_t<dot_kind::f16>;
+using kind_f8f6f4_t   = kind_t<dot_kind::f8f6f4>;
+using kind_i8_t       = kind_t<dot_kind::i8>;
+using kind_mxf4_t     = kind_t<dot_kind::mxf4>;
+using kind_mxf4nvf4_t = kind_t<dot_kind::mxf4nvf4>;
+using kind_mxf8f6f4_t = kind_t<dot_kind::mxf8f6f4>;
+using kind_tf32_t     = kind_t<dot_kind::tf32>;
+
+static constexpr kind_f16_t kind_f16{};
+static constexpr kind_f8f6f4_t kind_f8f6f4{};
+static constexpr kind_i8_t kind_i8{};
+static constexpr kind_mxf4_t kind_mxf4{};
+static constexpr kind_mxf4nvf4_t kind_mxf4nvf4{};
+static constexpr kind_mxf8f6f4_t kind_mxf8f6f4{};
+static constexpr kind_tf32_t kind_tf32{};
+
+template <int n>
+using n32_t = _CUDA_VSTD::integral_constant<int, n>;
+
 _LIBCUDACXX_END_NAMESPACE_CUDA_PTX
 
 #endif // _CUDA_PTX_DOT_VARIANTS_H_
diff --git a/libcudacxx/include/cuda/__ptx/ptx_helper_functions.h b/libcudacxx/include/cuda/__ptx/ptx_helper_functions.h
@@ -22,6 +22,7 @@
 #  pragma system_header
 #endif // no system header
 
+#include <cuda/std/__type_traits/enable_if.h>
 #include <cuda/std/__type_traits/integral_constant.h>
 #include <cuda/std/cstddef>
 #include <cuda/std/cstdint>
@@ -30,9 +31,6 @@
 
 _LIBCUDACXX_BEGIN_NAMESPACE_CUDA_PTX
 
-template <int __n>
-using n32_t = _CUDA_VSTD::integral_constant<int, __n>;
-
 /*************************************************************
  *
  * Conversion from generic pointer -> state space "pointer"