From e0c20a076f2ec84586b6508664df4f59273c6aaf Mon Sep 17 00:00:00 2001 From: "Harsh Aggarwal (NVIDIA)" Date: Wed, 20 Aug 2025 14:41:06 +0530 Subject: Updated support to enable batch3 (#8219) Enable CUDA support for batch 3 tests - Enhanced wave operations with exclusive support - Added proper identity values for min/max operations - Fixed intrinsic name mapping issues - Updated test configurations Co-authored-by: Ellie Hermaszewska --- prelude/slang-cuda-prelude.h | 295 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 293 insertions(+), 2 deletions(-) (limited to 'prelude') diff --git a/prelude/slang-cuda-prelude.h b/prelude/slang-cuda-prelude.h index 9df2727f6..44afd71b9 100644 --- a/prelude/slang-cuda-prelude.h +++ b/prelude/slang-cuda-prelude.h @@ -2661,17 +2661,67 @@ struct WaveOpMul template struct WaveOpMax { - __inline__ __device__ static T getInitial(T a) { return a; } + __inline__ __device__ static T getInitial(T a, bool exclusive = false); __inline__ __device__ static T doOp(T a, T b) { return a > b ? a : b; } }; template struct WaveOpMin { - __inline__ __device__ static T getInitial(T a) { return a; } + __inline__ __device__ static T getInitial(T a, bool exclusive = false); __inline__ __device__ static T doOp(T a, T b) { return a < b ? a : b; } }; +// Compact specializations using macro for getInitial +#define SLANG_WAVE_MIN_SPEC(T, EXCL_VAL) \ + template<> \ + __inline__ __device__ T WaveOpMin::getInitial(T a, bool exclusive) \ + { \ + return exclusive ? (EXCL_VAL) : a; \ + } + +#define SLANG_WAVE_MAX_SPEC(T, EXCL_VAL) \ + template<> \ + __inline__ __device__ T WaveOpMax::getInitial(T a, bool exclusive) \ + { \ + return exclusive ? (EXCL_VAL) : a; \ + } + +// Min specializations (exclusive identity = max value) +SLANG_WAVE_MIN_SPEC(float, SLANG_INFINITY) +SLANG_WAVE_MIN_SPEC(double, SLANG_INFINITY) +SLANG_WAVE_MIN_SPEC(int, 0x7FFFFFFF) +SLANG_WAVE_MIN_SPEC(uint, 0xFFFFFFFF) +SLANG_WAVE_MIN_SPEC(char, (char)0x7F) +SLANG_WAVE_MIN_SPEC(int8_t, (int8_t)0x7F) +SLANG_WAVE_MIN_SPEC(uint8_t, (uint8_t)0xFF) +SLANG_WAVE_MIN_SPEC(int16_t, (int16_t)0x7FFF) +SLANG_WAVE_MIN_SPEC(uint16_t, (uint16_t)0xFFFF) +SLANG_WAVE_MIN_SPEC(int64_t, 0x7FFFFFFFFFFFFFFFLL) +SLANG_WAVE_MIN_SPEC(uint64_t, 0xFFFFFFFFFFFFFFFFULL) +#if SLANG_CUDA_ENABLE_HALF +SLANG_WAVE_MIN_SPEC(__half, __ushort_as_half(0x7BFF)) +#endif + +// Max specializations (exclusive identity = min value) +SLANG_WAVE_MAX_SPEC(float, -SLANG_INFINITY) +SLANG_WAVE_MAX_SPEC(double, -SLANG_INFINITY) +SLANG_WAVE_MAX_SPEC(int, (int)0x80000000) +SLANG_WAVE_MAX_SPEC(uint, 0) +SLANG_WAVE_MAX_SPEC(char, (char)0x80) +SLANG_WAVE_MAX_SPEC(int8_t, (int8_t)0x80) +SLANG_WAVE_MAX_SPEC(uint8_t, 0) +SLANG_WAVE_MAX_SPEC(int16_t, (int16_t)0x8000) +SLANG_WAVE_MAX_SPEC(uint16_t, 0) +SLANG_WAVE_MAX_SPEC(int64_t, (int64_t)0x8000000000000000LL) +SLANG_WAVE_MAX_SPEC(uint64_t, 0) +#if SLANG_CUDA_ENABLE_HALF +SLANG_WAVE_MAX_SPEC(__half, __ushort_as_half(0xFBFF)) +#endif + +#undef SLANG_WAVE_MIN_SPEC +#undef SLANG_WAVE_MAX_SPEC + template struct ElementTypeTrait; @@ -2706,6 +2756,33 @@ struct ElementTypeTrait { typedef int64_t Type; }; +template<> +struct ElementTypeTrait +{ + typedef char Type; +}; +template<> +struct ElementTypeTrait +{ + typedef uchar Type; +}; +template<> +struct ElementTypeTrait +{ + typedef short Type; +}; +template<> +struct ElementTypeTrait +{ + typedef ushort Type; +}; +#if SLANG_CUDA_ENABLE_HALF +template<> +struct ElementTypeTrait<__half> +{ + typedef __half Type; +}; +#endif // Vector template<> @@ -2792,6 +2869,115 @@ struct ElementTypeTrait typedef double Type; }; +// Additional vector types +template<> +struct ElementTypeTrait +{ + typedef char Type; +}; +template<> +struct ElementTypeTrait +{ + typedef char Type; +}; +template<> +struct ElementTypeTrait +{ + typedef char Type; +}; +template<> +struct ElementTypeTrait +{ + typedef uchar Type; +}; +template<> +struct ElementTypeTrait +{ + typedef uchar Type; +}; +template<> +struct ElementTypeTrait +{ + typedef uchar Type; +}; +template<> +struct ElementTypeTrait +{ + typedef short Type; +}; +template<> +struct ElementTypeTrait +{ + typedef short Type; +}; +template<> +struct ElementTypeTrait +{ + typedef short Type; +}; +template<> +struct ElementTypeTrait +{ + typedef ushort Type; +}; +template<> +struct ElementTypeTrait +{ + typedef ushort Type; +}; +template<> +struct ElementTypeTrait +{ + typedef ushort Type; +}; +template<> +struct ElementTypeTrait +{ + typedef int64_t Type; +}; +template<> +struct ElementTypeTrait +{ + typedef int64_t Type; +}; +template<> +struct ElementTypeTrait +{ + typedef int64_t Type; +}; +template<> +struct ElementTypeTrait +{ + typedef uint64_t Type; +}; +template<> +struct ElementTypeTrait +{ + typedef uint64_t Type; +}; +template<> +struct ElementTypeTrait +{ + typedef uint64_t Type; +}; +#if SLANG_CUDA_ENABLE_HALF +template<> +struct ElementTypeTrait<__half2> +{ + typedef __half Type; +}; +template<> +struct ElementTypeTrait<__half3> +{ + typedef __half Type; +}; +template<> +struct ElementTypeTrait<__half4> +{ + typedef __half Type; +}; +#endif + // Matrix template struct ElementTypeTrait> @@ -3430,6 +3616,111 @@ __inline__ __device__ T _wavePrefixAndMultiple(WarpMask mask, T val) return val; } +template +__inline__ __device__ T _wavePrefixMin(WarpMask mask, T val) +{ + return _wavePrefixScalar, T>(mask, val); +} + +template +__inline__ __device__ T _wavePrefixMax(WarpMask mask, T val) +{ + return _wavePrefixScalar, T>(mask, val); +} + +template +__inline__ __device__ T _wavePrefixMinMultiple(WarpMask mask, T val) +{ + typedef typename ElementTypeTrait::Type ElemType; + _wavePrefixMultiple, ElemType, sizeof(T) / sizeof(ElemType)>( + mask, + (ElemType*)&val); + return val; +} + +template +__inline__ __device__ T _wavePrefixMaxMultiple(WarpMask mask, T val) +{ + typedef typename ElementTypeTrait::Type ElemType; + _wavePrefixMultiple, ElemType, sizeof(T) / sizeof(ElemType)>( + mask, + (ElemType*)&val); + return val; +} + +// Wrapper structures for exclusive operations that use the overloaded getInitial method +template +struct WaveOpExclusiveMin +{ + __inline__ __device__ static T getInitial(T a) { return WaveOpMin::getInitial(a, true); } + __inline__ __device__ static T doOp(T a, T b) { return WaveOpMin::doOp(a, b); } +}; + +template +struct WaveOpExclusiveMax +{ + __inline__ __device__ static T getInitial(T a) { return WaveOpMax::getInitial(a, true); } + __inline__ __device__ static T doOp(T a, T b) { return WaveOpMax::doOp(a, b); } +}; + +// Inclusive prefix min/max functions (for WaveMultiPrefixInclusive*) +template +__inline__ __device__ T _wavePrefixInclusiveMin(WarpMask mask, T val) +{ + return _wavePrefixMin(mask, val); +} + +template +__inline__ __device__ T _wavePrefixInclusiveMax(WarpMask mask, T val) +{ + return _wavePrefixMax(mask, val); +} + +template +__inline__ __device__ T _wavePrefixInclusiveMinMultiple(WarpMask mask, T val) +{ + return _wavePrefixMinMultiple(mask, val); +} + +template +__inline__ __device__ T _wavePrefixInclusiveMaxMultiple(WarpMask mask, T val) +{ + return _wavePrefixMaxMultiple(mask, val); +} + +// Explicit exclusive prefix min/max functions (for WaveMultiPrefixExclusive*) +template +__inline__ __device__ T _wavePrefixExclusiveMin(WarpMask mask, T val) +{ + return _wavePrefixScalar, T>(mask, val); +} + +template +__inline__ __device__ T _wavePrefixExclusiveMax(WarpMask mask, T val) +{ + return _wavePrefixScalar, T>(mask, val); +} + +template +__inline__ __device__ T _wavePrefixExclusiveMinMultiple(WarpMask mask, T val) +{ + typedef typename ElementTypeTrait::Type ElemType; + _wavePrefixMultiple, ElemType, sizeof(T) / sizeof(ElemType)>( + mask, + (ElemType*)&val); + return val; +} + +template +__inline__ __device__ T _wavePrefixExclusiveMaxMultiple(WarpMask mask, T val) +{ + typedef typename ElementTypeTrait::Type ElemType; + _wavePrefixMultiple, ElemType, sizeof(T) / sizeof(ElemType)>( + mask, + (ElemType*)&val); + return val; +} + template __inline__ __device__ uint4 _waveMatchScalar(WarpMask mask, T val) { -- cgit v1.2.3