yum-mirror/slang

Making it easier to work with shaders

git clone https://git.yummers.dev/yum-mirror/slang

Yong HeImmutable access qualifier for pointers and use `__ldg` on cuda. (#8710)01510f2c9

master
169.1 KiB4734 linesraw
1#define SLANG_PRELUDE_EXPORT
2
3#ifdef __CUDACC_RTC__
4#define SLANG_CUDA_RTC 1
5#else
6#define SLANG_CUDA_RTC 0
7#endif
8
9#if SLANG_CUDA_RTC
10
11#else
12
13#include <cstdint>
14#include <stdio.h>
15
16#endif
17
18// Define SLANG_CUDA_ENABLE_HALF to use the cuda_fp16 include to add half support.
19// For this to work NVRTC needs to have the path to the CUDA SDK.
20//
21// As it stands the includes paths defined for Slang are passed down to NVRTC. Similarly defines
22// defined for the Slang compile are passed down.
23
24#ifdef SLANG_CUDA_ENABLE_HALF
25// We don't want half2 operators, because it will implement comparison operators that return a
26// bool(!). We want to generate those functions. Doing so means that we will have to define all
27// the other half2 operators.
28#define __CUDA_NO_HALF2_OPERATORS__
29#include <cuda_fp16.h>
30#endif
31
32#ifdef SLANG_CUDA_ENABLE_OPTIX
33#include <optix.h>
34#endif
35
36// Define slang offsetof implementation
37#ifndef SLANG_OFFSET_OF
38#define SLANG_OFFSET_OF(type, member) (size_t)((char*)&(((type*)0)->member) - (char*)0)
39#endif
40
41// Must be large enough to cause overflow and therefore infinity
42#ifndef SLANG_INFINITY
43#define SLANG_INFINITY ((float)(1e+300 * 1e+300))
44#endif
45
46// For now we'll disable any asserts in this prelude
47#define SLANG_PRELUDE_ASSERT(x)
48
49#ifndef SLANG_CUDA_WARP_SIZE
50#define SLANG_CUDA_WARP_SIZE 32
51#endif
52
53#define SLANG_CUDA_WARP_MASK \
54    (SLANG_CUDA_WARP_SIZE - 1) // Used for masking threadIdx.x to the warp lane index
55#define SLANG_CUDA_WARP_BITMASK (~int(0))
56
57//
58#define SLANG_FORCE_INLINE inline
59
60#define SLANG_CUDA_CALL __device__
61
62#define SLANG_FORCE_INLINE inline
63#define SLANG_INLINE inline
64
65
66// Since we are using unsigned arithmatic care is need in this comparison.
67// It is *assumed* that sizeInBytes >= elemSize. Which means (sizeInBytes >= elemSize) >= 0
68// Which means only a single test is needed
69
70// Asserts for bounds checking.
71// It is assumed index/count are unsigned types.
72#define SLANG_BOUND_ASSERT(index, count) SLANG_PRELUDE_ASSERT(index < count);
73#define SLANG_BOUND_ASSERT_BYTE_ADDRESS(index, elemSize, sizeInBytes) \
74    SLANG_PRELUDE_ASSERT(index <= (sizeInBytes - elemSize) && (index & 3) == 0);
75
76// Macros to zero index if an access is out of range
77#define SLANG_BOUND_ZERO_INDEX(index, count) index = (index < count) ? index : 0;
78#define SLANG_BOUND_ZERO_INDEX_BYTE_ADDRESS(index, elemSize, sizeInBytes) \
79    index = (index <= (sizeInBytes - elemSize)) ? index : 0;
80
81// The 'FIX' macro define how the index is fixed. The default is to do nothing. If
82// SLANG_ENABLE_BOUND_ZERO_INDEX the fix macro will zero the index, if out of range
83#ifdef SLANG_ENABLE_BOUND_ZERO_INDEX
84#define SLANG_BOUND_FIX(index, count) SLANG_BOUND_ZERO_INDEX(index, count)
85#define SLANG_BOUND_FIX_BYTE_ADDRESS(index, elemSize, sizeInBytes) \
86    SLANG_BOUND_ZERO_INDEX_BYTE_ADDRESS(index, elemSize, sizeInBytes)
87#define SLANG_BOUND_FIX_FIXED_ARRAY(index, count) \
88    SLANG_BOUND_ZERO_INDEX(index, count) SLANG_BOUND_ZERO_INDEX(index, count)
89#else
90#define SLANG_BOUND_FIX(index, count)
91#define SLANG_BOUND_FIX_BYTE_ADDRESS(index, elemSize, sizeInBytes)
92#define SLANG_BOUND_FIX_FIXED_ARRAY(index, count)
93#endif
94
95#ifndef SLANG_BOUND_CHECK
96#define SLANG_BOUND_CHECK(index, count) \
97    SLANG_BOUND_ASSERT(index, count) SLANG_BOUND_FIX(index, count)
98#endif
99
100#ifndef SLANG_BOUND_CHECK_BYTE_ADDRESS
101#define SLANG_BOUND_CHECK_BYTE_ADDRESS(index, elemSize, sizeInBytes) \
102    SLANG_BOUND_ASSERT_BYTE_ADDRESS(index, elemSize, sizeInBytes)    \
103    SLANG_BOUND_FIX_BYTE_ADDRESS(index, elemSize, sizeInBytes)
104#endif
105
106#ifndef SLANG_BOUND_CHECK_FIXED_ARRAY
107#define SLANG_BOUND_CHECK_FIXED_ARRAY(index, count) \
108    SLANG_BOUND_ASSERT(index, count) SLANG_BOUND_FIX_FIXED_ARRAY(index, count)
109#endif
110
111// This macro handles how out-of-range surface coordinates are handled;
112// I can equal
113// cudaBoundaryModeClamp, in which case out-of-range coordinates are clamped to the valid range
114// cudaBoundaryModeZero, in which case out-of-range reads return zero and out-of-range writes are
115// ignored cudaBoundaryModeTrap, in which case out-of-range accesses cause the kernel execution to
116// fail.
117
118#ifndef SLANG_CUDA_BOUNDARY_MODE
119#define SLANG_CUDA_BOUNDARY_MODE cudaBoundaryModeZero
120
121// Can be one of SLANG_CUDA_PTX_BOUNDARY_MODE. Only applies *PTX* emitted CUDA operations
122// which currently is just RWTextureRW format writes
123//
124// .trap         causes an execution trap on out-of-bounds addresses
125// .clamp        stores data at the nearest surface location (sized appropriately)
126// .zero         drops stores to out-of-bounds addresses
127
128#define SLANG_PTX_BOUNDARY_MODE "zero"
129#endif
130
131struct TypeInfo
132{
133    size_t typeSize;
134};
135
136template<typename T, size_t SIZE>
137struct FixedArray
138{
139    SLANG_CUDA_CALL const T& operator[](size_t index) const
140    {
141        SLANG_BOUND_CHECK_FIXED_ARRAY(index, SIZE);
142        return m_data[index];
143    }
144    SLANG_CUDA_CALL T& operator[](size_t index)
145    {
146        SLANG_BOUND_CHECK_FIXED_ARRAY(index, SIZE);
147        return m_data[index];
148    }
149
150    T m_data[SIZE];
151};
152
153// An array that has no specified size, becomes a 'Array'. This stores the size so it can
154// potentially do bounds checking.
155template<typename T>
156struct Array
157{
158    SLANG_CUDA_CALL const T& operator[](size_t index) const
159    {
160        SLANG_BOUND_CHECK(index, count);
161        return data[index];
162    }
163    SLANG_CUDA_CALL T& operator[](size_t index)
164    {
165        SLANG_BOUND_CHECK(index, count);
166        return data[index];
167    }
168
169    T* data;
170    size_t count;
171};
172
173// Typically defined in cuda.h, but we can't ship/rely on that, so just define here
174typedef unsigned long long CUtexObject;
175typedef unsigned long long CUsurfObject;
176
177// On CUDA sampler state is actually bound up with the texture object. We have a SamplerState type,
178// backed as a pointer, to simplify code generation, with the downside that such a binding will take
179// up uniform space, even though it will have no effect.
180// TODO(JS): Consider ways to strip use of variables of this type so have no binding,
181struct SamplerStateUnused;
182typedef SamplerStateUnused* SamplerState;
183
184
185// TODO(JS): Not clear yet if this can be handled on CUDA, by just ignoring.
186// For now, just map to the index type.
187typedef size_t NonUniformResourceIndex;
188
189// Code generator will generate the specific type
190template<typename T, int ROWS, int COLS>
191struct Matrix;
192
193// Boolean vector types should follow CUDA's builtin vector alignment rules
194// Align boolX the same as charX according to CUDA spec:
195// char1/uchar1: 1-byte aligned, char2/uchar2: 2-byte aligned
196// char3/uchar3: 1-byte aligned, char4/uchar4: 4-byte aligned
197struct __align__(1) bool1
198{
199    bool x;
200
201    SLANG_FORCE_INLINE SLANG_CUDA_CALL bool& operator[](int idx)
202    {
203        return (&x)[idx];
204    }
205    SLANG_FORCE_INLINE SLANG_CUDA_CALL const bool& operator[](int idx) const
206    {
207        return (&x)[idx];
208    }
209};
210
211struct __align__(2) bool2
212{
213    bool x, y;
214
215    SLANG_FORCE_INLINE SLANG_CUDA_CALL bool& operator[](int idx)
216    {
217        return (&x)[idx];
218    }
219    SLANG_FORCE_INLINE SLANG_CUDA_CALL const bool& operator[](int idx) const
220    {
221        return (&x)[idx];
222    }
223};
224
225struct __align__(1) bool3
226{
227    bool x, y, z;
228
229    SLANG_FORCE_INLINE SLANG_CUDA_CALL bool& operator[](int idx)
230    {
231        return (&x)[idx];
232    }
233    SLANG_FORCE_INLINE SLANG_CUDA_CALL const bool& operator[](int idx) const
234    {
235        return (&x)[idx];
236    }
237};
238
239struct __align__(4) bool4
240{
241    bool x, y, z, w;
242
243    SLANG_FORCE_INLINE SLANG_CUDA_CALL bool& operator[](int idx)
244    {
245        return (&x)[idx];
246    }
247    SLANG_FORCE_INLINE SLANG_CUDA_CALL const bool& operator[](int idx) const
248    {
249        return (&x)[idx];
250    }
251};
252
253SLANG_FORCE_INLINE SLANG_CUDA_CALL bool __ldg(const bool* ptr)
254{
255    return (bool)(__ldg((const char*)ptr));
256}
257
258SLANG_FORCE_INLINE SLANG_CUDA_CALL bool2 __ldg(const bool2* ptr)
259{
260    auto val = __ldg((const char2*)ptr);
261    return {val.x != 0, val.y != 0};
262}
263
264SLANG_FORCE_INLINE SLANG_CUDA_CALL bool4 __ldg(const bool4* ptr)
265{
266    auto val = __ldg((const char4*)ptr);
267    return {val.x != 0, val.y != 0, val.z != 0, val.w != 0};
268}
269
270#if SLANG_CUDA_RTC
271
272typedef signed char int8_t;
273typedef short int16_t;
274typedef int int32_t;
275typedef long long int64_t;
276typedef ptrdiff_t intptr_t;
277
278typedef unsigned char uint8_t;
279typedef unsigned short uint16_t;
280typedef unsigned int uint32_t;
281typedef unsigned long long uint64_t;
282typedef size_t uintptr_t;
283
284#endif
285
286typedef long long longlong;
287typedef unsigned long long ulonglong;
288
289typedef unsigned char uchar;
290typedef unsigned short ushort;
291typedef unsigned int uint;
292
293union Union32
294{
295    uint32_t u;
296    int32_t i;
297    float f;
298};
299
300union Union64
301{
302    uint64_t u;
303    int64_t i;
304    double d;
305};
306
307template<typename T>
308SLANG_FORCE_INLINE SLANG_CUDA_CALL float make_float(T val)
309{
310    return (float)val;
311}
312
313SLANG_FORCE_INLINE SLANG_CUDA_CALL float _slang_fmod(float x, float y)
314{
315    return ::fmodf(x, y);
316}
317SLANG_FORCE_INLINE SLANG_CUDA_CALL double _slang_fmod(double x, double y)
318{
319    return ::fmod(x, y);
320}
321
322#if SLANG_CUDA_ENABLE_HALF
323
324// Add the other vector half types
325struct __half1
326{
327    __half x;
328};
329struct __align__(4) __half3
330{
331    __half x, y, z;
332};
333struct __align__(4) __half4
334{
335    __half x, y, z, w;
336};
337#endif
338
339#define SLANG_VECTOR_GET_ELEMENT(T)                                                   \
340    SLANG_FORCE_INLINE SLANG_CUDA_CALL T _slang_vector_get_element(T##1 x, int index) \
341    {                                                                                 \
342        return ((T*)(&x))[index];                                                     \
343    }                                                                                 \
344    SLANG_FORCE_INLINE SLANG_CUDA_CALL T _slang_vector_get_element(T##2 x, int index) \
345    {                                                                                 \
346        return ((T*)(&x))[index];                                                     \
347    }                                                                                 \
348    SLANG_FORCE_INLINE SLANG_CUDA_CALL T _slang_vector_get_element(T##3 x, int index) \
349    {                                                                                 \
350        return ((T*)(&x))[index];                                                     \
351    }                                                                                 \
352    SLANG_FORCE_INLINE SLANG_CUDA_CALL T _slang_vector_get_element(T##4 x, int index) \
353    {                                                                                 \
354        return ((T*)(&x))[index];                                                     \
355    }
356SLANG_VECTOR_GET_ELEMENT(int)
357SLANG_VECTOR_GET_ELEMENT(bool)
358SLANG_VECTOR_GET_ELEMENT(uint)
359SLANG_VECTOR_GET_ELEMENT(short)
360SLANG_VECTOR_GET_ELEMENT(ushort)
361SLANG_VECTOR_GET_ELEMENT(char)
362SLANG_VECTOR_GET_ELEMENT(uchar)
363SLANG_VECTOR_GET_ELEMENT(longlong)
364SLANG_VECTOR_GET_ELEMENT(ulonglong)
365SLANG_VECTOR_GET_ELEMENT(float)
366SLANG_VECTOR_GET_ELEMENT(double)
367
368#define SLANG_VECTOR_GET_ELEMENT_PTR(T)                                                            \
369    SLANG_FORCE_INLINE SLANG_CUDA_CALL T* _slang_vector_get_element_ptr(const T##1 * x, int index) \
370    {                                                                                              \
371        return ((T*)(x)) + index;                                                                  \
372    }                                                                                              \
373    SLANG_FORCE_INLINE SLANG_CUDA_CALL T* _slang_vector_get_element_ptr(const T##2 * x, int index) \
374    {                                                                                              \
375        return ((T*)(x)) + index;                                                                  \
376    }                                                                                              \
377    SLANG_FORCE_INLINE SLANG_CUDA_CALL T* _slang_vector_get_element_ptr(const T##3 * x, int index) \
378    {                                                                                              \
379        return ((T*)(x)) + index;                                                                  \
380    }                                                                                              \
381    SLANG_FORCE_INLINE SLANG_CUDA_CALL T* _slang_vector_get_element_ptr(const T##4 * x, int index) \
382    {                                                                                              \
383        return ((T*)(x)) + index;                                                                  \
384    }
385SLANG_VECTOR_GET_ELEMENT_PTR(int)
386SLANG_VECTOR_GET_ELEMENT_PTR(bool)
387SLANG_VECTOR_GET_ELEMENT_PTR(uint)
388SLANG_VECTOR_GET_ELEMENT_PTR(short)
389SLANG_VECTOR_GET_ELEMENT_PTR(ushort)
390SLANG_VECTOR_GET_ELEMENT_PTR(char)
391SLANG_VECTOR_GET_ELEMENT_PTR(uchar)
392SLANG_VECTOR_GET_ELEMENT_PTR(longlong)
393SLANG_VECTOR_GET_ELEMENT_PTR(ulonglong)
394SLANG_VECTOR_GET_ELEMENT_PTR(float)
395SLANG_VECTOR_GET_ELEMENT_PTR(double)
396
397#if SLANG_CUDA_ENABLE_HALF
398SLANG_VECTOR_GET_ELEMENT(__half)
399SLANG_VECTOR_GET_ELEMENT_PTR(__half)
400#endif
401
402#define SLANG_CUDA_VECTOR_BINARY_OP(T, n, op)                                                 \
403    SLANG_FORCE_INLINE SLANG_CUDA_CALL T##n operator op(T##n thisVal, T##n other)             \
404    {                                                                                         \
405        T##n result;                                                                          \
406        for (int i = 0; i < n; i++)                                                           \
407            *_slang_vector_get_element_ptr(&result, i) =                                      \
408                _slang_vector_get_element(thisVal, i) op _slang_vector_get_element(other, i); \
409        return result;                                                                        \
410    }
411#define SLANG_CUDA_VECTOR_BINARY_COMPARE_OP(T, n, op)                                           \
412    SLANG_FORCE_INLINE SLANG_CUDA_CALL bool##n operator op(T##n thisVal, T##n other)            \
413    {                                                                                           \
414        bool##n result;                                                                         \
415        for (int i = 0; i < n; i++)                                                             \
416            *_slang_vector_get_element_ptr(&result, i) =                                        \
417                (_slang_vector_get_element(thisVal, i) op _slang_vector_get_element(other, i)); \
418        return result;                                                                          \
419    }
420#define SLANG_CUDA_VECTOR_UNARY_OP(T, n, op)                                                       \
421    SLANG_FORCE_INLINE SLANG_CUDA_CALL T##n operator op(T##n thisVal)                              \
422    {                                                                                              \
423        T##n result;                                                                               \
424        for (int i = 0; i < n; i++)                                                                \
425            *_slang_vector_get_element_ptr(&result, i) = op _slang_vector_get_element(thisVal, i); \
426        return result;                                                                             \
427    }
428
429#define SLANG_CUDA_VECTOR_INT_OP(T, n)            \
430    SLANG_CUDA_VECTOR_BINARY_OP(T, n, +)          \
431    SLANG_CUDA_VECTOR_BINARY_OP(T, n, -)          \
432    SLANG_CUDA_VECTOR_BINARY_OP(T, n, *)          \
433    SLANG_CUDA_VECTOR_BINARY_OP(T, n, /)          \
434    SLANG_CUDA_VECTOR_BINARY_OP(T, n, %)          \
435    SLANG_CUDA_VECTOR_BINARY_OP(T, n, ^)          \
436    SLANG_CUDA_VECTOR_BINARY_OP(T, n, &)          \
437    SLANG_CUDA_VECTOR_BINARY_OP(T, n, |)          \
438    SLANG_CUDA_VECTOR_BINARY_OP(T, n, &&)         \
439    SLANG_CUDA_VECTOR_BINARY_OP(T, n, ||)         \
440    SLANG_CUDA_VECTOR_BINARY_OP(T, n, >>)         \
441    SLANG_CUDA_VECTOR_BINARY_OP(T, n, <<)         \
442    SLANG_CUDA_VECTOR_BINARY_COMPARE_OP(T, n, >)  \
443    SLANG_CUDA_VECTOR_BINARY_COMPARE_OP(T, n, <)  \
444    SLANG_CUDA_VECTOR_BINARY_COMPARE_OP(T, n, >=) \
445    SLANG_CUDA_VECTOR_BINARY_COMPARE_OP(T, n, <=) \
446    SLANG_CUDA_VECTOR_BINARY_COMPARE_OP(T, n, ==) \
447    SLANG_CUDA_VECTOR_BINARY_COMPARE_OP(T, n, !=) \
448    SLANG_CUDA_VECTOR_UNARY_OP(T, n, !)           \
449    SLANG_CUDA_VECTOR_UNARY_OP(T, n, -)           \
450    SLANG_CUDA_VECTOR_UNARY_OP(T, n, ~)
451
452#define SLANG_CUDA_VECTOR_INT_OPS(T) \
453    SLANG_CUDA_VECTOR_INT_OP(T, 2)   \
454    SLANG_CUDA_VECTOR_INT_OP(T, 3)   \
455    SLANG_CUDA_VECTOR_INT_OP(T, 4)
456
457SLANG_CUDA_VECTOR_INT_OPS(int)
458SLANG_CUDA_VECTOR_INT_OPS(bool)
459SLANG_CUDA_VECTOR_INT_OPS(uint)
460SLANG_CUDA_VECTOR_INT_OPS(ushort)
461SLANG_CUDA_VECTOR_INT_OPS(short)
462SLANG_CUDA_VECTOR_INT_OPS(char)
463SLANG_CUDA_VECTOR_INT_OPS(uchar)
464SLANG_CUDA_VECTOR_INT_OPS(longlong)
465SLANG_CUDA_VECTOR_INT_OPS(ulonglong)
466
467#define SLANG_CUDA_VECTOR_FLOAT_OP(T, n)          \
468    SLANG_CUDA_VECTOR_BINARY_OP(T, n, +)          \
469    SLANG_CUDA_VECTOR_BINARY_OP(T, n, -)          \
470    SLANG_CUDA_VECTOR_BINARY_OP(T, n, *)          \
471    SLANG_CUDA_VECTOR_BINARY_OP(T, n, /)          \
472    SLANG_CUDA_VECTOR_BINARY_OP(T, n, &&)         \
473    SLANG_CUDA_VECTOR_BINARY_OP(T, n, ||)         \
474    SLANG_CUDA_VECTOR_BINARY_COMPARE_OP(T, n, >)  \
475    SLANG_CUDA_VECTOR_BINARY_COMPARE_OP(T, n, <)  \
476    SLANG_CUDA_VECTOR_BINARY_COMPARE_OP(T, n, >=) \
477    SLANG_CUDA_VECTOR_BINARY_COMPARE_OP(T, n, <=) \
478    SLANG_CUDA_VECTOR_BINARY_COMPARE_OP(T, n, ==) \
479    SLANG_CUDA_VECTOR_BINARY_COMPARE_OP(T, n, !=) \
480    SLANG_CUDA_VECTOR_UNARY_OP(T, n, -)
481#define SLANG_CUDA_VECTOR_FLOAT_OPS(T) \
482    SLANG_CUDA_VECTOR_FLOAT_OP(T, 2)   \
483    SLANG_CUDA_VECTOR_FLOAT_OP(T, 3)   \
484    SLANG_CUDA_VECTOR_FLOAT_OP(T, 4)
485
486SLANG_CUDA_VECTOR_FLOAT_OPS(float)
487SLANG_CUDA_VECTOR_FLOAT_OPS(double)
488#if SLANG_CUDA_ENABLE_HALF
489SLANG_CUDA_VECTOR_FLOAT_OPS(__half)
490#endif
491#define SLANG_CUDA_FLOAT_VECTOR_MOD_IMPL(T, n)                                             \
492    SLANG_FORCE_INLINE SLANG_CUDA_CALL T##n operator%(const T##n& left, const T##n& right) \
493    {                                                                                      \
494        T##n result;                                                                       \
495        for (int i = 0; i < n; i++)                                                        \
496            *_slang_vector_get_element_ptr(&result, i) = _slang_fmod(                      \
497                _slang_vector_get_element(left, i),                                        \
498                _slang_vector_get_element(right, i));                                      \
499        return result;                                                                     \
500    }
501#define SLANG_CUDA_FLOAT_VECTOR_MOD(T)     \
502    SLANG_CUDA_FLOAT_VECTOR_MOD_IMPL(T, 2) \
503    SLANG_CUDA_FLOAT_VECTOR_MOD_IMPL(T, 3) \
504    SLANG_CUDA_FLOAT_VECTOR_MOD_IMPL(T, 4)
505
506SLANG_CUDA_FLOAT_VECTOR_MOD(float)
507SLANG_CUDA_FLOAT_VECTOR_MOD(double)
508
509#if SLANG_CUDA_RTC || SLANG_CUDA_ENABLE_HALF
510#define SLANG_MAKE_VECTOR(T)                                                \
511    SLANG_FORCE_INLINE SLANG_CUDA_CALL T##2 make_##T##2(T x, T y)           \
512    {                                                                       \
513        return T##2 {x, y};                                                 \
514    }                                                                       \
515    SLANG_FORCE_INLINE SLANG_CUDA_CALL T##3 make_##T##3(T x, T y, T z)      \
516    {                                                                       \
517        return T##3 {x, y, z};                                              \
518    }                                                                       \
519    SLANG_FORCE_INLINE SLANG_CUDA_CALL T##4 make_##T##4(T x, T y, T z, T w) \
520    {                                                                       \
521        return T##4 {x, y, z, w};                                           \
522    }
523#endif
524
525#if SLANG_CUDA_RTC
526SLANG_MAKE_VECTOR(int)
527SLANG_MAKE_VECTOR(uint)
528SLANG_MAKE_VECTOR(short)
529SLANG_MAKE_VECTOR(ushort)
530SLANG_MAKE_VECTOR(char)
531SLANG_MAKE_VECTOR(uchar)
532SLANG_MAKE_VECTOR(float)
533SLANG_MAKE_VECTOR(double)
534SLANG_MAKE_VECTOR(longlong)
535SLANG_MAKE_VECTOR(ulonglong)
536#endif
537
538#if SLANG_CUDA_ENABLE_HALF
539SLANG_MAKE_VECTOR(__half)
540#endif
541
542SLANG_FORCE_INLINE SLANG_CUDA_CALL bool1 make_bool1(bool x)
543{
544    return bool1{x};
545}
546SLANG_FORCE_INLINE SLANG_CUDA_CALL bool2 make_bool2(bool x, bool y)
547{
548    return bool2{x, y};
549}
550SLANG_FORCE_INLINE SLANG_CUDA_CALL bool3 make_bool3(bool x, bool y, bool z)
551{
552    return bool3{x, y, z};
553}
554SLANG_FORCE_INLINE SLANG_CUDA_CALL bool4 make_bool4(bool x, bool y, bool z, bool w)
555{
556    return bool4{x, y, z, w};
557}
558SLANG_FORCE_INLINE SLANG_CUDA_CALL bool2 make_bool2(bool x)
559{
560    return bool2{x, x};
561}
562SLANG_FORCE_INLINE SLANG_CUDA_CALL bool3 make_bool3(bool x)
563{
564    return bool3{x, x, x};
565}
566SLANG_FORCE_INLINE SLANG_CUDA_CALL bool4 make_bool4(bool x)
567{
568    return bool4{x, x, x, x};
569}
570
571#if SLANG_CUDA_RTC
572#define SLANG_MAKE_VECTOR_FROM_SCALAR(T)                     \
573    SLANG_FORCE_INLINE SLANG_CUDA_CALL T##1 make_##T##1(T x) \
574    {                                                        \
575        return T##1 {x};                                     \
576    }                                                        \
577    SLANG_FORCE_INLINE SLANG_CUDA_CALL T##2 make_##T##2(T x) \
578    {                                                        \
579        return make_##T##2(x, x);                            \
580    }                                                        \
581    SLANG_FORCE_INLINE SLANG_CUDA_CALL T##3 make_##T##3(T x) \
582    {                                                        \
583        return make_##T##3(x, x, x);                         \
584    }                                                        \
585    SLANG_FORCE_INLINE SLANG_CUDA_CALL T##4 make_##T##4(T x) \
586    {                                                        \
587        return make_##T##4(x, x, x, x);                      \
588    }
589#else
590#define SLANG_MAKE_VECTOR_FROM_SCALAR(T)                     \
591    SLANG_FORCE_INLINE SLANG_CUDA_CALL T##2 make_##T##2(T x) \
592    {                                                        \
593        return make_##T##2(x, x);                            \
594    }                                                        \
595    SLANG_FORCE_INLINE SLANG_CUDA_CALL T##3 make_##T##3(T x) \
596    {                                                        \
597        return make_##T##3(x, x, x);                         \
598    }                                                        \
599    SLANG_FORCE_INLINE SLANG_CUDA_CALL T##4 make_##T##4(T x) \
600    {                                                        \
601        return make_##T##4(x, x, x, x);                      \
602    }
603#endif
604SLANG_MAKE_VECTOR_FROM_SCALAR(int)
605SLANG_MAKE_VECTOR_FROM_SCALAR(uint)
606SLANG_MAKE_VECTOR_FROM_SCALAR(short)
607SLANG_MAKE_VECTOR_FROM_SCALAR(ushort)
608SLANG_MAKE_VECTOR_FROM_SCALAR(char)
609SLANG_MAKE_VECTOR_FROM_SCALAR(uchar)
610SLANG_MAKE_VECTOR_FROM_SCALAR(longlong)
611SLANG_MAKE_VECTOR_FROM_SCALAR(ulonglong)
612SLANG_MAKE_VECTOR_FROM_SCALAR(float)
613SLANG_MAKE_VECTOR_FROM_SCALAR(double)
614#if SLANG_CUDA_ENABLE_HALF
615SLANG_MAKE_VECTOR_FROM_SCALAR(__half)
616#if !SLANG_CUDA_RTC
617SLANG_FORCE_INLINE SLANG_CUDA_CALL __half1 make___half1(__half x)
618{
619    return __half1{x};
620}
621#endif
622#endif
623
624#define SLANG_CUDA_VECTOR_ATOMIC_BINARY_IMPL(Fn, T, N)                                            \
625    SLANG_FORCE_INLINE SLANG_CUDA_CALL T##N Fn(T##N* address, T##N val)                           \
626    {                                                                                             \
627        T##N result;                                                                              \
628        for (int i = 0; i < N; i++)                                                               \
629            *_slang_vector_get_element_ptr(&result, i) =                                          \
630                Fn(_slang_vector_get_element_ptr(address, i), _slang_vector_get_element(val, i)); \
631        return result;                                                                            \
632    }
633
634#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 900
635SLANG_CUDA_VECTOR_ATOMIC_BINARY_IMPL(atomicAdd, float, 2)
636SLANG_CUDA_VECTOR_ATOMIC_BINARY_IMPL(atomicAdd, float, 4)
637#endif
638SLANG_CUDA_VECTOR_ATOMIC_BINARY_IMPL(atomicAdd, float, 3)
639SLANG_CUDA_VECTOR_ATOMIC_BINARY_IMPL(atomicAdd, int, 2)
640SLANG_CUDA_VECTOR_ATOMIC_BINARY_IMPL(atomicAdd, int, 3)
641SLANG_CUDA_VECTOR_ATOMIC_BINARY_IMPL(atomicAdd, int, 4)
642SLANG_CUDA_VECTOR_ATOMIC_BINARY_IMPL(atomicAdd, uint, 2)
643SLANG_CUDA_VECTOR_ATOMIC_BINARY_IMPL(atomicAdd, uint, 3)
644SLANG_CUDA_VECTOR_ATOMIC_BINARY_IMPL(atomicAdd, uint, 4)
645SLANG_CUDA_VECTOR_ATOMIC_BINARY_IMPL(atomicAdd, ulonglong, 2)
646SLANG_CUDA_VECTOR_ATOMIC_BINARY_IMPL(atomicAdd, ulonglong, 3)
647SLANG_CUDA_VECTOR_ATOMIC_BINARY_IMPL(atomicAdd, ulonglong, 4)
648
649template<typename T, int n>
650struct GetVectorTypeImpl
651{
652};
653
654#define GET_VECTOR_TYPE_IMPL(T, n)                                     \
655    template<>                                                         \
656    struct GetVectorTypeImpl<T, n>                                     \
657    {                                                                  \
658        typedef T##n type;                                             \
659        static SLANG_FORCE_INLINE SLANG_CUDA_CALL T##n fromScalar(T v) \
660        {                                                              \
661            return make_##T##n(v);                                     \
662        }                                                              \
663    };
664#define GET_VECTOR_TYPE_IMPL_N(T) \
665    GET_VECTOR_TYPE_IMPL(T, 1)    \
666    GET_VECTOR_TYPE_IMPL(T, 2)    \
667    GET_VECTOR_TYPE_IMPL(T, 3)    \
668    GET_VECTOR_TYPE_IMPL(T, 4)
669
670GET_VECTOR_TYPE_IMPL_N(int)
671GET_VECTOR_TYPE_IMPL_N(bool)
672GET_VECTOR_TYPE_IMPL_N(uint)
673GET_VECTOR_TYPE_IMPL_N(short)
674GET_VECTOR_TYPE_IMPL_N(ushort)
675GET_VECTOR_TYPE_IMPL_N(char)
676GET_VECTOR_TYPE_IMPL_N(uchar)
677GET_VECTOR_TYPE_IMPL_N(longlong)
678GET_VECTOR_TYPE_IMPL_N(ulonglong)
679GET_VECTOR_TYPE_IMPL_N(float)
680GET_VECTOR_TYPE_IMPL_N(double)
681#if SLANG_CUDA_ENABLE_HALF
682GET_VECTOR_TYPE_IMPL_N(__half)
683#endif
684template<typename T, int n>
685using Vector = typename GetVectorTypeImpl<T, n>::type;
686
687template<typename T, int n, typename OtherT, int m>
688SLANG_FORCE_INLINE SLANG_CUDA_CALL Vector<T, n> _slang_vector_reshape(const Vector<OtherT, m> other)
689{
690    Vector<T, n> result;
691    for (int i = 0; i < n; i++)
692    {
693        OtherT otherElement = T(0);
694        if (i < m)
695            otherElement = _slang_vector_get_element(other, i);
696        *_slang_vector_get_element_ptr(&result, i) = (T)otherElement;
697    }
698    return result;
699}
700
701template<typename T, int ROWS, int COLS>
702struct Matrix
703{
704    Vector<T, COLS> rows[ROWS];
705    SLANG_FORCE_INLINE SLANG_CUDA_CALL Vector<T, COLS>& operator[](size_t index)
706    {
707        return rows[index];
708    }
709
710    SLANG_FORCE_INLINE SLANG_CUDA_CALL const Vector<T, COLS>& operator[](size_t index) const
711    {
712        return rows[index];
713    }
714};
715
716
717template<typename T, int ROWS, int COLS>
718SLANG_FORCE_INLINE SLANG_CUDA_CALL Matrix<T, ROWS, COLS> makeMatrix(T scalar)
719{
720    Matrix<T, ROWS, COLS> result;
721    for (int i = 0; i < ROWS; i++)
722        result.rows[i] = GetVectorTypeImpl<T, COLS>::fromScalar(scalar);
723    return result;
724}
725
726template<typename T, int ROWS, int COLS>
727SLANG_FORCE_INLINE SLANG_CUDA_CALL Matrix<T, ROWS, COLS> makeMatrix(const Vector<T, COLS>& row0)
728{
729    Matrix<T, ROWS, COLS> result;
730    result.rows[0] = row0;
731    return result;
732}
733
734template<typename T, int ROWS, int COLS>
735SLANG_FORCE_INLINE SLANG_CUDA_CALL Matrix<T, ROWS, COLS> makeMatrix(
736    const Vector<T, COLS>& row0,
737    const Vector<T, COLS>& row1)
738{
739    Matrix<T, ROWS, COLS> result;
740    result.rows[0] = row0;
741    result.rows[1] = row1;
742    return result;
743}
744
745template<typename T, int ROWS, int COLS>
746SLANG_FORCE_INLINE SLANG_CUDA_CALL Matrix<T, ROWS, COLS> makeMatrix(
747    const Vector<T, COLS>& row0,
748    const Vector<T, COLS>& row1,
749    const Vector<T, COLS>& row2)
750{
751    Matrix<T, ROWS, COLS> result;
752    result.rows[0] = row0;
753    result.rows[1] = row1;
754    result.rows[2] = row2;
755    return result;
756}
757
758template<typename T, int ROWS, int COLS>
759SLANG_FORCE_INLINE SLANG_CUDA_CALL Matrix<T, ROWS, COLS> makeMatrix(
760    const Vector<T, COLS>& row0,
761    const Vector<T, COLS>& row1,
762    const Vector<T, COLS>& row2,
763    const Vector<T, COLS>& row3)
764{
765    Matrix<T, ROWS, COLS> result;
766    result.rows[0] = row0;
767    result.rows[1] = row1;
768    result.rows[2] = row2;
769    result.rows[3] = row3;
770    return result;
771}
772
773template<typename T, int ROWS, int COLS, typename U, int otherRow, int otherCol>
774SLANG_FORCE_INLINE SLANG_CUDA_CALL Matrix<T, ROWS, COLS> makeMatrix(
775    const Matrix<U, otherRow, otherCol>& other)
776{
777    Matrix<T, ROWS, COLS> result;
778    int minRow = ROWS;
779    int minCol = COLS;
780    if (minRow > otherRow)
781        minRow = otherRow;
782    if (minCol > otherCol)
783        minCol = otherCol;
784    for (int i = 0; i < minRow; i++)
785        for (int j = 0; j < minCol; j++)
786            *_slang_vector_get_element_ptr(result.rows + i, j) =
787                (T)_slang_vector_get_element(other.rows[i], j);
788    return result;
789}
790
791template<typename T, int ROWS, int COLS>
792SLANG_FORCE_INLINE SLANG_CUDA_CALL Matrix<T, ROWS, COLS> makeMatrix(T v0, T v1, T v2, T v3)
793{
794    Matrix<T, ROWS, COLS> rs;
795    rs.rows[0].x = v0;
796    rs.rows[0].y = v1;
797    rs.rows[1].x = v2;
798    rs.rows[1].y = v3;
799    return rs;
800}
801
802template<typename T, int ROWS, int COLS>
803SLANG_FORCE_INLINE SLANG_CUDA_CALL Matrix<T, ROWS, COLS> makeMatrix(
804    T v0,
805    T v1,
806    T v2,
807    T v3,
808    T v4,
809    T v5)
810{
811    Matrix<T, ROWS, COLS> rs;
812    if (COLS == 3)
813    {
814        *_slang_vector_get_element_ptr(&rs.rows[0], 0) = v0;
815        *_slang_vector_get_element_ptr(&rs.rows[0], 1) = v1;
816        *_slang_vector_get_element_ptr(&rs.rows[0], 2) = v2;
817        *_slang_vector_get_element_ptr(&rs.rows[1], 0) = v3;
818        *_slang_vector_get_element_ptr(&rs.rows[1], 1) = v4;
819        *_slang_vector_get_element_ptr(&rs.rows[1], 2) = v5;
820    }
821    else
822    {
823        rs.rows[0].x = v0;
824        rs.rows[0].y = v1;
825        rs.rows[1].x = v2;
826        rs.rows[1].y = v3;
827        rs.rows[2].x = v4;
828        rs.rows[2].y = v5;
829    }
830    return rs;
831}
832
833template<typename T, int ROWS, int COLS>
834SLANG_FORCE_INLINE SLANG_CUDA_CALL Matrix<T, ROWS, COLS> makeMatrix(
835    T v0,
836    T v1,
837    T v2,
838    T v3,
839    T v4,
840    T v5,
841    T v6,
842    T v7)
843{
844    Matrix<T, ROWS, COLS> rs;
845    if (COLS == 4)
846    {
847        *_slang_vector_get_element_ptr(&rs.rows[0], 0) = v0;
848        *_slang_vector_get_element_ptr(&rs.rows[0], 1) = v1;
849        *_slang_vector_get_element_ptr(&rs.rows[0], 2) = v2;
850        *_slang_vector_get_element_ptr(&rs.rows[0], 3) = v3;
851        *_slang_vector_get_element_ptr(&rs.rows[1], 0) = v4;
852        *_slang_vector_get_element_ptr(&rs.rows[1], 1) = v5;
853        *_slang_vector_get_element_ptr(&rs.rows[1], 2) = v6;
854        *_slang_vector_get_element_ptr(&rs.rows[1], 3) = v7;
855    }
856    else
857    {
858        rs.rows[0].x = v0;
859        rs.rows[0].y = v1;
860        rs.rows[1].x = v2;
861        rs.rows[1].y = v3;
862        rs.rows[2].x = v4;
863        rs.rows[2].y = v5;
864        rs.rows[3].x = v6;
865        rs.rows[3].y = v7;
866    }
867    return rs;
868}
869
870template<typename T, int ROWS, int COLS>
871SLANG_FORCE_INLINE SLANG_CUDA_CALL Matrix<T, ROWS, COLS> makeMatrix(
872    T v0,
873    T v1,
874    T v2,
875    T v3,
876    T v4,
877    T v5,
878    T v6,
879    T v7,
880    T v8)
881{
882    Matrix<T, ROWS, COLS> rs;
883    rs.rows[0].x = v0;
884    rs.rows[0].y = v1;
885    rs.rows[0].z = v2;
886    rs.rows[1].x = v3;
887    rs.rows[1].y = v4;
888    rs.rows[1].z = v5;
889    rs.rows[2].x = v6;
890    rs.rows[2].y = v7;
891    rs.rows[2].z = v8;
892    return rs;
893}
894
895template<typename T, int ROWS, int COLS>
896SLANG_FORCE_INLINE SLANG_CUDA_CALL Matrix<T, ROWS, COLS> makeMatrix(
897    T v0,
898    T v1,
899    T v2,
900    T v3,
901    T v4,
902    T v5,
903    T v6,
904    T v7,
905    T v8,
906    T v9,
907    T v10,
908    T v11)
909{
910    Matrix<T, ROWS, COLS> rs;
911    if (COLS == 4)
912    {
913        *_slang_vector_get_element_ptr(&rs.rows[0], 0) = v0;
914        *_slang_vector_get_element_ptr(&rs.rows[0], 1) = v1;
915        *_slang_vector_get_element_ptr(&rs.rows[0], 2) = v2;
916        *_slang_vector_get_element_ptr(&rs.rows[0], 3) = v3;
917        *_slang_vector_get_element_ptr(&rs.rows[1], 0) = v4;
918        *_slang_vector_get_element_ptr(&rs.rows[1], 1) = v5;
919        *_slang_vector_get_element_ptr(&rs.rows[1], 2) = v6;
920        *_slang_vector_get_element_ptr(&rs.rows[1], 3) = v7;
921        *_slang_vector_get_element_ptr(&rs.rows[2], 0) = v8;
922        *_slang_vector_get_element_ptr(&rs.rows[2], 1) = v9;
923        *_slang_vector_get_element_ptr(&rs.rows[2], 2) = v10;
924        *_slang_vector_get_element_ptr(&rs.rows[2], 3) = v11;
925    }
926    else
927    {
928        rs.rows[0].x = v0;
929        rs.rows[0].y = v1;
930        rs.rows[0].z = v2;
931        rs.rows[1].x = v3;
932        rs.rows[1].y = v4;
933        rs.rows[1].z = v5;
934        rs.rows[2].x = v6;
935        rs.rows[2].y = v7;
936        rs.rows[2].z = v8;
937        rs.rows[3].x = v9;
938        rs.rows[3].y = v10;
939        rs.rows[3].z = v11;
940    }
941    return rs;
942}
943
944template<typename T, int ROWS, int COLS>
945SLANG_FORCE_INLINE SLANG_CUDA_CALL Matrix<T, ROWS, COLS> makeMatrix(
946    T v0,
947    T v1,
948    T v2,
949    T v3,
950    T v4,
951    T v5,
952    T v6,
953    T v7,
954    T v8,
955    T v9,
956    T v10,
957    T v11,
958    T v12,
959    T v13,
960    T v14,
961    T v15)
962{
963    Matrix<T, ROWS, COLS> rs;
964    rs.rows[0].x = v0;
965    rs.rows[0].y = v1;
966    rs.rows[0].z = v2;
967    rs.rows[0].w = v3;
968    rs.rows[1].x = v4;
969    rs.rows[1].y = v5;
970    rs.rows[1].z = v6;
971    rs.rows[1].w = v7;
972    rs.rows[2].x = v8;
973    rs.rows[2].y = v9;
974    rs.rows[2].z = v10;
975    rs.rows[2].w = v11;
976    rs.rows[3].x = v12;
977    rs.rows[3].y = v13;
978    rs.rows[3].z = v14;
979    rs.rows[3].w = v15;
980    return rs;
981}
982
983#define SLANG_MATRIX_BINARY_OP(T, op)                                   \
984    template<int R, int C>                                              \
985    SLANG_FORCE_INLINE SLANG_CUDA_CALL Matrix<T, R, C> operator op(     \
986        const Matrix<T, R, C>& thisVal,                                 \
987        const Matrix<T, R, C>& other)                                   \
988    {                                                                   \
989        Matrix<T, R, C> result;                                         \
990        for (int i = 0; i < R; i++)                                     \
991            for (int j = 0; j < C; j++)                                 \
992                *_slang_vector_get_element_ptr(result.rows + i, j) =    \
993                    _slang_vector_get_element(thisVal.rows[i], j)       \
994                        op _slang_vector_get_element(other.rows[i], j); \
995        return result;                                                  \
996    }
997
998#define SLANG_MATRIX_UNARY_OP(T, op)                                                               \
999    template<int R, int C>                                                                         \
1000    SLANG_FORCE_INLINE SLANG_CUDA_CALL Matrix<T, R, C> operator op(const Matrix<T, R, C>& thisVal) \
1001    {                                                                                              \
1002        Matrix<T, R, C> result;                                                                    \
1003        for (int i = 0; i < R; i++)                                                                \
1004            for (int j = 0; j < C; j++)                                                            \
1005                *_slang_vector_get_element_ptr(result.rows + i, j) =                               \
1006                    op _slang_vector_get_element(thisVal.rows[i], j);                              \
1007        return result;                                                                             \
1008    }
1009#define SLANG_INT_MATRIX_OPS(T)   \
1010    SLANG_MATRIX_BINARY_OP(T, +)  \
1011    SLANG_MATRIX_BINARY_OP(T, -)  \
1012    SLANG_MATRIX_BINARY_OP(T, *)  \
1013    SLANG_MATRIX_BINARY_OP(T, /)  \
1014    SLANG_MATRIX_BINARY_OP(T, &)  \
1015    SLANG_MATRIX_BINARY_OP(T, |)  \
1016    SLANG_MATRIX_BINARY_OP(T, &&) \
1017    SLANG_MATRIX_BINARY_OP(T, ||) \
1018    SLANG_MATRIX_BINARY_OP(T, ^)  \
1019    SLANG_MATRIX_BINARY_OP(T, %)  \
1020    SLANG_MATRIX_UNARY_OP(T, !)   \
1021    SLANG_MATRIX_UNARY_OP(T, ~)
1022#define SLANG_FLOAT_MATRIX_OPS(T) \
1023    SLANG_MATRIX_BINARY_OP(T, +)  \
1024    SLANG_MATRIX_BINARY_OP(T, -)  \
1025    SLANG_MATRIX_BINARY_OP(T, *)  \
1026    SLANG_MATRIX_BINARY_OP(T, /)  \
1027    SLANG_MATRIX_UNARY_OP(T, -)
1028SLANG_INT_MATRIX_OPS(int)
1029SLANG_INT_MATRIX_OPS(uint)
1030SLANG_INT_MATRIX_OPS(short)
1031SLANG_INT_MATRIX_OPS(ushort)
1032SLANG_INT_MATRIX_OPS(char)
1033SLANG_INT_MATRIX_OPS(uchar)
1034SLANG_INT_MATRIX_OPS(longlong)
1035SLANG_INT_MATRIX_OPS(ulonglong)
1036SLANG_FLOAT_MATRIX_OPS(float)
1037SLANG_FLOAT_MATRIX_OPS(double)
1038#if SLANG_CUDA_ENABLE_HALF
1039SLANG_FLOAT_MATRIX_OPS(__half)
1040#endif
1041#define SLANG_MATRIX_INT_NEG_OP(T)                                                        \
1042    template<int R, int C>                                                                \
1043    SLANG_FORCE_INLINE SLANG_CUDA_CALL Matrix<T, R, C> operator-(Matrix<T, R, C> thisVal) \
1044    {                                                                                     \
1045        Matrix<T, R, C> result;                                                           \
1046        for (int i = 0; i < R; i++)                                                       \
1047            for (int j = 0; j < C; j++)                                                   \
1048                *_slang_vector_get_element_ptr(result.rows + i, j) =                      \
1049                    0 - _slang_vector_get_element(thisVal.rows[i], j);                    \
1050        return result;                                                                    \
1051    }
1052SLANG_MATRIX_INT_NEG_OP(int)
1053SLANG_MATRIX_INT_NEG_OP(uint)
1054SLANG_MATRIX_INT_NEG_OP(short)
1055SLANG_MATRIX_INT_NEG_OP(ushort)
1056SLANG_MATRIX_INT_NEG_OP(char)
1057SLANG_MATRIX_INT_NEG_OP(uchar)
1058SLANG_MATRIX_INT_NEG_OP(longlong)
1059SLANG_MATRIX_INT_NEG_OP(ulonglong)
1060
1061#define SLANG_FLOAT_MATRIX_MOD(T)                                                 \
1062    template<int R, int C>                                                        \
1063    SLANG_FORCE_INLINE SLANG_CUDA_CALL Matrix<T, R, C> operator%(                 \
1064        Matrix<T, R, C> left,                                                     \
1065        Matrix<T, R, C> right)                                                    \
1066    {                                                                             \
1067        Matrix<T, R, C> result;                                                   \
1068        for (int i = 0; i < R; i++)                                               \
1069            for (int j = 0; j < C; j++)                                           \
1070                *_slang_vector_get_element_ptr(result.rows + i, j) = _slang_fmod( \
1071                    _slang_vector_get_element(left.rows[i], j),                   \
1072                    _slang_vector_get_element(right.rows[i], j));                 \
1073        return result;                                                            \
1074    }
1075
1076SLANG_FLOAT_MATRIX_MOD(float)
1077SLANG_FLOAT_MATRIX_MOD(double)
1078#if SLANG_CUDA_ENABLE_HALF
1079template<int R, int C>
1080SLANG_FORCE_INLINE SLANG_CUDA_CALL Matrix<__half, R, C> operator%(
1081    Matrix<__half, R, C> left,
1082    Matrix<__half, R, C> right)
1083{
1084    Matrix<__half, R, C> result;
1085    for (int i = 0; i < R; i++)
1086        for (int j = 0; j < C; j++)
1087            *_slang_vector_get_element_ptr(result.rows + i, j) = __float2half(_slang_fmod(
1088                __half2float(_slang_vector_get_element(left.rows[i], j)),
1089                __half2float(_slang_vector_get_element(right.rows[i], j))));
1090    return result;
1091}
1092#endif
1093#undef SLANG_FLOAT_MATRIX_MOD
1094#undef SLANG_MATRIX_BINARY_OP
1095#undef SLANG_MATRIX_UNARY_OP
1096#undef SLANG_INT_MATRIX_OPS
1097#undef SLANG_FLOAT_MATRIX_OPS
1098#undef SLANG_MATRIX_INT_NEG_OP
1099#undef SLANG_FLOAT_MATRIX_MOD
1100
1101#define SLANG_SELECT_IMPL(T, N)                                                                  \
1102    SLANG_FORCE_INLINE SLANG_CUDA_CALL Vector<T, N> _slang_select(                               \
1103        bool##N condition,                                                                       \
1104        Vector<T, N> v0,                                                                         \
1105        Vector<T, N> v1)                                                                         \
1106    {                                                                                            \
1107        Vector<T, N> result;                                                                     \
1108        for (int i = 0; i < N; i++)                                                              \
1109        {                                                                                        \
1110            *_slang_vector_get_element_ptr(&result, i) = _slang_vector_get_element(condition, i) \
1111                                                             ? _slang_vector_get_element(v0, i)  \
1112                                                             : _slang_vector_get_element(v1, i); \
1113        }                                                                                        \
1114        return result;                                                                           \
1115    }
1116#define SLANG_SELECT_T(T)   \
1117    SLANG_SELECT_IMPL(T, 2) \
1118    SLANG_SELECT_IMPL(T, 3) \
1119    SLANG_SELECT_IMPL(T, 4)
1120
1121SLANG_SELECT_T(int)
1122SLANG_SELECT_T(bool)
1123SLANG_SELECT_T(uint)
1124SLANG_SELECT_T(short)
1125SLANG_SELECT_T(ushort)
1126SLANG_SELECT_T(char)
1127SLANG_SELECT_T(uchar)
1128SLANG_SELECT_T(float)
1129SLANG_SELECT_T(double)
1130
1131template<typename T>
1132SLANG_FORCE_INLINE SLANG_CUDA_CALL T _slang_select(bool condition, T v0, T v1)
1133{
1134    return condition ? v0 : v1;
1135}
1136
1137//
1138// Half support
1139//
1140
1141#if SLANG_CUDA_ENABLE_HALF
1142SLANG_SELECT_T(__half)
1143
1144// Convenience functions ushort -> half
1145
1146SLANG_FORCE_INLINE SLANG_CUDA_CALL __half2 __ushort_as_half(const ushort2& i)
1147{
1148    return __halves2half2(__ushort_as_half(i.x), __ushort_as_half(i.y));
1149}
1150SLANG_FORCE_INLINE SLANG_CUDA_CALL __half3 __ushort_as_half(const ushort3& i)
1151{
1152    return __half3{__ushort_as_half(i.x), __ushort_as_half(i.y), __ushort_as_half(i.z)};
1153}
1154SLANG_FORCE_INLINE SLANG_CUDA_CALL __half4 __ushort_as_half(const ushort4& i)
1155{
1156    return __half4{
1157        __ushort_as_half(i.x),
1158        __ushort_as_half(i.y),
1159        __ushort_as_half(i.z),
1160        __ushort_as_half(i.w)};
1161}
1162
1163// Convenience functions half -> ushort
1164
1165SLANG_FORCE_INLINE SLANG_CUDA_CALL ushort2 __half_as_ushort(const __half2& i)
1166{
1167    return make_ushort2(__half_as_ushort(i.x), __half_as_ushort(i.y));
1168}
1169SLANG_FORCE_INLINE SLANG_CUDA_CALL ushort3 __half_as_ushort(const __half3& i)
1170{
1171    return make_ushort3(__half_as_ushort(i.x), __half_as_ushort(i.y), __half_as_ushort(i.z));
1172}
1173SLANG_FORCE_INLINE SLANG_CUDA_CALL ushort4 __half_as_ushort(const __half4& i)
1174{
1175    return make_ushort4(
1176        __half_as_ushort(i.x),
1177        __half_as_ushort(i.y),
1178        __half_as_ushort(i.z),
1179        __half_as_ushort(i.w));
1180}
1181
1182// This is a little bit of a hack. Fortunately CUDA has the definitions of the templated types in
1183// include/surface_indirect_functions.h
1184// Here we find the template definition requires a specialization of __nv_isurf_trait to allow
1185// a specialization of the surface write functions.
1186// This *isn't* a problem on the read functions as they don't have a return type that uses this
1187// mechanism
1188
1189template<>
1190struct __nv_isurf_trait<__half>
1191{
1192    typedef void type;
1193};
1194template<>
1195struct __nv_isurf_trait<__half2>
1196{
1197    typedef void type;
1198};
1199template<>
1200struct __nv_isurf_trait<__half4>
1201{
1202    typedef void type;
1203};
1204
1205#define SLANG_DROP_PARENS(...) __VA_ARGS__
1206
1207#define SLANG_SURFACE_READ(FUNC_NAME, TYPE_ARGS, ARGS)                                             \
1208    template<>                                                                                     \
1209    SLANG_FORCE_INLINE SLANG_CUDA_CALL __half FUNC_NAME<__half>(                                   \
1210        cudaSurfaceObject_t surfObj,                                                               \
1211        SLANG_DROP_PARENS TYPE_ARGS,                                                               \
1212        cudaSurfaceBoundaryMode boundaryMode)                                                      \
1213    {                                                                                              \
1214        return __ushort_as_half(FUNC_NAME<ushort>(surfObj, SLANG_DROP_PARENS ARGS, boundaryMode)); \
1215    }                                                                                              \
1216                                                                                                   \
1217    template<>                                                                                     \
1218    SLANG_FORCE_INLINE SLANG_CUDA_CALL __half2 FUNC_NAME<__half2>(                                 \
1219        cudaSurfaceObject_t surfObj,                                                               \
1220        SLANG_DROP_PARENS TYPE_ARGS,                                                               \
1221        cudaSurfaceBoundaryMode boundaryMode)                                                      \
1222    {                                                                                              \
1223        return __ushort_as_half(                                                                   \
1224            FUNC_NAME<ushort2>(surfObj, SLANG_DROP_PARENS ARGS, boundaryMode));                    \
1225    }                                                                                              \
1226                                                                                                   \
1227    template<>                                                                                     \
1228    SLANG_FORCE_INLINE SLANG_CUDA_CALL __half4 FUNC_NAME<__half4>(                                 \
1229        cudaSurfaceObject_t surfObj,                                                               \
1230        SLANG_DROP_PARENS TYPE_ARGS,                                                               \
1231        cudaSurfaceBoundaryMode boundaryMode)                                                      \
1232    {                                                                                              \
1233        return __ushort_as_half(                                                                   \
1234            FUNC_NAME<ushort4>(surfObj, SLANG_DROP_PARENS ARGS, boundaryMode));                    \
1235    }
1236
1237SLANG_SURFACE_READ(surf1Dread, (int x), (x))
1238SLANG_SURFACE_READ(surf2Dread, (int x, int y), (x, y))
1239SLANG_SURFACE_READ(surf3Dread, (int x, int y, int z), (x, y, z))
1240SLANG_SURFACE_READ(surf1DLayeredread, (int x, int layer), (x, layer))
1241SLANG_SURFACE_READ(surf2DLayeredread, (int x, int y, int layer), (x, y, layer))
1242SLANG_SURFACE_READ(surfCubemapread, (int x, int y, int face), (x, y, face))
1243SLANG_SURFACE_READ(surfCubemapLayeredread, (int x, int y, int layerFace), (x, y, layerFace))
1244
1245#define SLANG_SURFACE_WRITE(FUNC_NAME, TYPE_ARGS, ARGS)                                            \
1246    template<>                                                                                     \
1247    SLANG_FORCE_INLINE SLANG_CUDA_CALL void FUNC_NAME<__half>(                                     \
1248        __half data,                                                                               \
1249        cudaSurfaceObject_t surfObj,                                                               \
1250        SLANG_DROP_PARENS TYPE_ARGS,                                                               \
1251        cudaSurfaceBoundaryMode boundaryMode)                                                      \
1252    {                                                                                              \
1253        FUNC_NAME<ushort>(__half_as_ushort(data), surfObj, SLANG_DROP_PARENS ARGS, boundaryMode);  \
1254    }                                                                                              \
1255                                                                                                   \
1256    template<>                                                                                     \
1257    SLANG_FORCE_INLINE SLANG_CUDA_CALL void FUNC_NAME<__half2>(                                    \
1258        __half2 data,                                                                              \
1259        cudaSurfaceObject_t surfObj,                                                               \
1260        SLANG_DROP_PARENS TYPE_ARGS,                                                               \
1261        cudaSurfaceBoundaryMode boundaryMode)                                                      \
1262    {                                                                                              \
1263        FUNC_NAME<ushort2>(__half_as_ushort(data), surfObj, SLANG_DROP_PARENS ARGS, boundaryMode); \
1264    }                                                                                              \
1265                                                                                                   \
1266    template<>                                                                                     \
1267    SLANG_FORCE_INLINE SLANG_CUDA_CALL void FUNC_NAME<__half4>(                                    \
1268        __half4 data,                                                                              \
1269        cudaSurfaceObject_t surfObj,                                                               \
1270        SLANG_DROP_PARENS TYPE_ARGS,                                                               \
1271        cudaSurfaceBoundaryMode boundaryMode)                                                      \
1272    {                                                                                              \
1273        FUNC_NAME<ushort4>(__half_as_ushort(data), surfObj, SLANG_DROP_PARENS ARGS, boundaryMode); \
1274    }
1275
1276SLANG_SURFACE_WRITE(surf1Dwrite, (int x), (x))
1277SLANG_SURFACE_WRITE(surf2Dwrite, (int x, int y), (x, y))
1278SLANG_SURFACE_WRITE(surf3Dwrite, (int x, int y, int z), (x, y, z))
1279SLANG_SURFACE_WRITE(surf1DLayeredwrite, (int x, int layer), (x, layer))
1280SLANG_SURFACE_WRITE(surf2DLayeredwrite, (int x, int y, int layer), (x, y, layer))
1281SLANG_SURFACE_WRITE(surfCubemapwrite, (int x, int y, int face), (x, y, face))
1282SLANG_SURFACE_WRITE(surfCubemapLayeredwrite, (int x, int y, int layerFace), (x, y, layerFace))
1283
1284// ! Hack to test out reading !!!
1285// Only works converting *from* half
1286
1287// template <typename T>
1288// SLANG_FORCE_INLINE SLANG_CUDA_CALL T surf2Dread_convert(cudaSurfaceObject_t surfObj, int x, int
1289// y, cudaSurfaceBoundaryMode boundaryMode);
1290
1291#define SLANG_SURFACE_READ_HALF_CONVERT(FUNC_NAME, TYPE_ARGS, ARGS)                              \
1292                                                                                                 \
1293    template<typename T>                                                                         \
1294    SLANG_FORCE_INLINE SLANG_CUDA_CALL T FUNC_NAME##_convert(                                    \
1295        cudaSurfaceObject_t surfObj,                                                             \
1296        SLANG_DROP_PARENS TYPE_ARGS,                                                             \
1297        cudaSurfaceBoundaryMode boundaryMode);                                                   \
1298                                                                                                 \
1299    template<>                                                                                   \
1300    SLANG_FORCE_INLINE SLANG_CUDA_CALL float FUNC_NAME##_convert<float>(                         \
1301        cudaSurfaceObject_t surfObj,                                                             \
1302        SLANG_DROP_PARENS TYPE_ARGS,                                                             \
1303        cudaSurfaceBoundaryMode boundaryMode)                                                    \
1304    {                                                                                            \
1305        return __ushort_as_half(                                                                 \
1306            FUNC_NAME<uint16_t>(surfObj, SLANG_DROP_PARENS ARGS, boundaryMode));                 \
1307    }                                                                                            \
1308                                                                                                 \
1309    template<>                                                                                   \
1310    SLANG_FORCE_INLINE SLANG_CUDA_CALL float2 FUNC_NAME##_convert<float2>(                       \
1311        cudaSurfaceObject_t surfObj,                                                             \
1312        SLANG_DROP_PARENS TYPE_ARGS,                                                             \
1313        cudaSurfaceBoundaryMode boundaryMode)                                                    \
1314    {                                                                                            \
1315        const __half2 v =                                                                        \
1316            __ushort_as_half(FUNC_NAME<ushort2>(surfObj, SLANG_DROP_PARENS ARGS, boundaryMode)); \
1317        return float2{v.x, v.y};                                                                 \
1318    }                                                                                            \
1319                                                                                                 \
1320    template<>                                                                                   \
1321    SLANG_FORCE_INLINE SLANG_CUDA_CALL float4 FUNC_NAME##_convert<float4>(                       \
1322        cudaSurfaceObject_t surfObj,                                                             \
1323        SLANG_DROP_PARENS TYPE_ARGS,                                                             \
1324        cudaSurfaceBoundaryMode boundaryMode)                                                    \
1325    {                                                                                            \
1326        const __half4 v =                                                                        \
1327            __ushort_as_half(FUNC_NAME<ushort4>(surfObj, SLANG_DROP_PARENS ARGS, boundaryMode)); \
1328        return float4{v.x, v.y, v.z, v.w};                                                       \
1329    }
1330
1331SLANG_SURFACE_READ_HALF_CONVERT(surf1Dread, (int x), (x))
1332SLANG_SURFACE_READ_HALF_CONVERT(surf2Dread, (int x, int y), (x, y))
1333SLANG_SURFACE_READ_HALF_CONVERT(surf3Dread, (int x, int y, int z), (x, y, z))
1334
1335#endif
1336
1337// Support for doing format conversion when writing to a surface/RWTexture
1338
1339// NOTE! For normal surface access x values are *byte* addressed.
1340// For the _convert versions they are *not*. They don't need to be because sust.p does not require
1341// it.
1342
1343// https://docs.nvidia.com/cuda/inline-ptx-assembly/index.html
1344// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#surface-instructions-sust
1345
1346
1347// surf1Dwrite_convert
1348
1349template<typename T>
1350SLANG_FORCE_INLINE SLANG_CUDA_CALL void surf1Dwrite_convert(
1351    T v,
1352    cudaSurfaceObject_t surfObj,
1353    int x,
1354    cudaSurfaceBoundaryMode boundaryMode);
1355
1356#define SLANG_SURF1DWRITE_CONVERT_IMPL(T, c)                                                     \
1357    template<>                                                                                   \
1358    SLANG_FORCE_INLINE SLANG_CUDA_CALL void surf1Dwrite_convert<T>(                              \
1359        T v,                                                                                     \
1360        cudaSurfaceObject_t surfObj,                                                             \
1361        int x,                                                                                   \
1362        cudaSurfaceBoundaryMode boundaryMode)                                                    \
1363    {                                                                                            \
1364        asm volatile(                                                                            \
1365            "sust.p.1d.b32." SLANG_PTX_BOUNDARY_MODE " [%0, {%1}], {%2};" ::"l"(surfObj),        \
1366            "r"(x),                                                                              \
1367            c(v));                                                                               \
1368    }                                                                                            \
1369    template<>                                                                                   \
1370    SLANG_FORCE_INLINE SLANG_CUDA_CALL void surf1Dwrite_convert<T##2>(                           \
1371        T##2 v,                                                                                  \
1372        cudaSurfaceObject_t surfObj,                                                             \
1373        int x,                                                                                   \
1374        cudaSurfaceBoundaryMode boundaryMode)                                                    \
1375    {                                                                                            \
1376        const T vx = v.x, vy = v.y;                                                              \
1377        asm volatile(                                                                            \
1378            "sust.p.1d.v2.b32." SLANG_PTX_BOUNDARY_MODE " [%0, {%1}], {%2, %3};" ::"l"(surfObj), \
1379            "r"(x),                                                                              \
1380            c(vx),                                                                               \
1381            c(vy));                                                                              \
1382    }                                                                                            \
1383    template<>                                                                                   \
1384    SLANG_FORCE_INLINE SLANG_CUDA_CALL void surf1Dwrite_convert<T##4>(                           \
1385        T##4 v,                                                                                  \
1386        cudaSurfaceObject_t surfObj,                                                             \
1387        int x,                                                                                   \
1388        cudaSurfaceBoundaryMode boundaryMode)                                                    \
1389    {                                                                                            \
1390        const T vx = v.x, vy = v.y, vz = v.z, vw = v.w;                                          \
1391        asm volatile(                                                                            \
1392            "sust.p.1d.v4.b32." SLANG_PTX_BOUNDARY_MODE                                          \
1393            " [%0, {%1}], {%2, %3, %4, %5};" ::"l"(surfObj),                                     \
1394            "r"(x),                                                                              \
1395            c(vx),                                                                               \
1396            c(vy),                                                                               \
1397            c(vz),                                                                               \
1398            c(vw));                                                                              \
1399    }
1400
1401SLANG_SURF1DWRITE_CONVERT_IMPL(float, "f")
1402SLANG_SURF1DWRITE_CONVERT_IMPL(uint, "r")
1403SLANG_SURF1DWRITE_CONVERT_IMPL(int, "r")
1404
1405// surf1DLayeredwrite_convert (not supported)
1406
1407template<typename T>
1408SLANG_FORCE_INLINE SLANG_CUDA_CALL void surf1DLayeredwrite_convert(
1409    T v,
1410    cudaSurfaceObject_t surfObj,
1411    int x,
1412    int layer,
1413    cudaSurfaceBoundaryMode boundaryMode)
1414{
1415    static_assert(false, "CUDA doesn't support formatted surface writes on 1D array surfaces");
1416}
1417
1418// surf2Dwrite_convert
1419
1420template<typename T>
1421SLANG_FORCE_INLINE SLANG_CUDA_CALL void surf2Dwrite_convert(
1422    T v,
1423    cudaSurfaceObject_t surfObj,
1424    int x,
1425    int y,
1426    cudaSurfaceBoundaryMode boundaryMode);
1427
1428#define SLANG_SURF2DWRITE_CONVERT_IMPL(T, c)                                                  \
1429    template<>                                                                                \
1430    SLANG_FORCE_INLINE SLANG_CUDA_CALL void surf2Dwrite_convert<T>(                           \
1431        T v,                                                                                  \
1432        cudaSurfaceObject_t surfObj,                                                          \
1433        int x,                                                                                \
1434        int y,                                                                                \
1435        cudaSurfaceBoundaryMode boundaryMode)                                                 \
1436    {                                                                                         \
1437        asm volatile(                                                                         \
1438            "sust.p.2d.b32." SLANG_PTX_BOUNDARY_MODE " [%0, {%1, %2}], {%3};" ::"l"(surfObj), \
1439            "r"(x),                                                                           \
1440            "r"(y),                                                                           \
1441            c(v));                                                                            \
1442    }                                                                                         \
1443    template<>                                                                                \
1444    SLANG_FORCE_INLINE SLANG_CUDA_CALL void surf2Dwrite_convert<T##2>(                        \
1445        T##2 v,                                                                               \
1446        cudaSurfaceObject_t surfObj,                                                          \
1447        int x,                                                                                \
1448        int y,                                                                                \
1449        cudaSurfaceBoundaryMode boundaryMode)                                                 \
1450    {                                                                                         \
1451        const T vx = v.x, vy = v.y;                                                           \
1452        asm volatile(                                                                         \
1453            "sust.p.2d.v2.b32." SLANG_PTX_BOUNDARY_MODE                                       \
1454            " [%0, {%1, %2}], {%3, %4};" ::"l"(surfObj),                                      \
1455            "r"(x),                                                                           \
1456            "r"(y),                                                                           \
1457            c(vx),                                                                            \
1458            c(vy));                                                                           \
1459    }                                                                                         \
1460    template<>                                                                                \
1461    SLANG_FORCE_INLINE SLANG_CUDA_CALL void surf2Dwrite_convert<T##4>(                        \
1462        T##4 v,                                                                               \
1463        cudaSurfaceObject_t surfObj,                                                          \
1464        int x,                                                                                \
1465        int y,                                                                                \
1466        cudaSurfaceBoundaryMode boundaryMode)                                                 \
1467    {                                                                                         \
1468        const T vx = v.x, vy = v.y, vz = v.z, vw = v.w;                                       \
1469        asm volatile(                                                                         \
1470            "sust.p.2d.v4.b32." SLANG_PTX_BOUNDARY_MODE                                       \
1471            " [%0, {%1, %2}], {%3, %4, %5, %6};" ::"l"(surfObj),                              \
1472            "r"(x),                                                                           \
1473            "r"(y),                                                                           \
1474            c(vx),                                                                            \
1475            c(vy),                                                                            \
1476            c(vz),                                                                            \
1477            c(vw));                                                                           \
1478    }
1479
1480SLANG_SURF2DWRITE_CONVERT_IMPL(float, "f")
1481SLANG_SURF2DWRITE_CONVERT_IMPL(uint, "r")
1482SLANG_SURF2DWRITE_CONVERT_IMPL(int, "r")
1483
1484// surf2DLayeredwrite_convert (not supported)
1485
1486template<typename T>
1487SLANG_FORCE_INLINE SLANG_CUDA_CALL void surf2DLayeredwrite_convert(
1488    T v,
1489    cudaSurfaceObject_t surfObj,
1490    int x,
1491    int y,
1492    int layer,
1493    cudaSurfaceBoundaryMode boundaryMode)
1494{
1495    static_assert(false, "CUDA doesn't support formatted surface writes on 2D array surfaces");
1496}
1497
1498// surf3Dwrite_convert
1499
1500template<typename T>
1501SLANG_FORCE_INLINE SLANG_CUDA_CALL void surf3Dwrite_convert(
1502    T v,
1503    cudaSurfaceObject_t surfObj,
1504    int x,
1505    int y,
1506    int z,
1507    cudaSurfaceBoundaryMode boundaryMode);
1508
1509#define SLANG_SURF3DWRITE_CONVERT_IMPL(T, c)                             \
1510    template<>                                                           \
1511    SLANG_FORCE_INLINE SLANG_CUDA_CALL void surf3Dwrite_convert<T>(      \
1512        T v,                                                             \
1513        cudaSurfaceObject_t surfObj,                                     \
1514        int x,                                                           \
1515        int y,                                                           \
1516        int z,                                                           \
1517        cudaSurfaceBoundaryMode boundaryMode)                            \
1518    {                                                                    \
1519        asm volatile(                                                    \
1520            "sust.p.3d.b32." SLANG_PTX_BOUNDARY_MODE                     \
1521            " [%0, {%1, %2, %3, %4}], {%5};" ::"l"(surfObj),             \
1522            "r"(x),                                                      \
1523            "r"(y),                                                      \
1524            "r"(z),                                                      \
1525            "r"(0),                                                      \
1526            c(v));                                                       \
1527    }                                                                    \
1528    template<>                                                           \
1529    SLANG_FORCE_INLINE SLANG_CUDA_CALL void surf3Dwrite_convert<T##2>(   \
1530        T##2 v,                                                          \
1531        cudaSurfaceObject_t surfObj,                                     \
1532        int x,                                                           \
1533        int y,                                                           \
1534        int z,                                                           \
1535        cudaSurfaceBoundaryMode boundaryMode)                            \
1536    {                                                                    \
1537        const T vx = v.x, vy = v.y;                                      \
1538        asm volatile(                                                    \
1539            "sust.p.3d.v2.b32." SLANG_PTX_BOUNDARY_MODE                  \
1540            " [%0, {%1, %2, %3, %4}], {%5, %6};" ::"l"(surfObj),         \
1541            "r"(x),                                                      \
1542            "r"(y),                                                      \
1543            "r"(z),                                                      \
1544            "r"(0),                                                      \
1545            c(vx),                                                       \
1546            c(vy));                                                      \
1547    }                                                                    \
1548    template<>                                                           \
1549    SLANG_FORCE_INLINE SLANG_CUDA_CALL void surf3Dwrite_convert<T##4>(   \
1550        T##4 v,                                                          \
1551        cudaSurfaceObject_t surfObj,                                     \
1552        int x,                                                           \
1553        int y,                                                           \
1554        int z,                                                           \
1555        cudaSurfaceBoundaryMode boundaryMode)                            \
1556    {                                                                    \
1557        const T vx = v.x, vy = v.y, vz = v.z, vw = v.w;                  \
1558        asm volatile(                                                    \
1559            "sust.p.3d.v4.b32." SLANG_PTX_BOUNDARY_MODE                  \
1560            " [%0, {%1, %2, %3, %4}], {%5, %6, %7, %8};" ::"l"(surfObj), \
1561            "r"(x),                                                      \
1562            "r"(y),                                                      \
1563            "r"(z),                                                      \
1564            "r"(0),                                                      \
1565            c(vx),                                                       \
1566            c(vy),                                                       \
1567            c(vz),                                                       \
1568            c(vw));                                                      \
1569    }
1570
1571SLANG_SURF3DWRITE_CONVERT_IMPL(float, "f")
1572SLANG_SURF3DWRITE_CONVERT_IMPL(uint, "r")
1573SLANG_SURF3DWRITE_CONVERT_IMPL(int, "r")
1574
1575// ----------------------------- F32 -----------------------------------------
1576
1577// Unary
1578SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_ceil(float f)
1579{
1580    return ::ceilf(f);
1581}
1582SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_floor(float f)
1583{
1584    return ::floorf(f);
1585}
1586SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_round(float f)
1587{
1588    return ::roundf(f);
1589}
1590SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_sin(float f)
1591{
1592    return ::sinf(f);
1593}
1594SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_cos(float f)
1595{
1596    return ::cosf(f);
1597}
1598SLANG_FORCE_INLINE SLANG_CUDA_CALL void F32_sincos(float f, float* s, float* c)
1599{
1600    ::sincosf(f, s, c);
1601}
1602SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_tan(float f)
1603{
1604    return ::tanf(f);
1605}
1606SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_asin(float f)
1607{
1608    return ::asinf(f);
1609}
1610SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_acos(float f)
1611{
1612    return ::acosf(f);
1613}
1614SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_atan(float f)
1615{
1616    return ::atanf(f);
1617}
1618SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_sinh(float f)
1619{
1620    return ::sinhf(f);
1621}
1622SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_cosh(float f)
1623{
1624    return ::coshf(f);
1625}
1626SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_tanh(float f)
1627{
1628    return ::tanhf(f);
1629}
1630SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_asinh(float f)
1631{
1632    return ::asinhf(f);
1633}
1634SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_acosh(float f)
1635{
1636    return ::acoshf(f);
1637}
1638SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_atanh(float f)
1639{
1640    return ::atanhf(f);
1641}
1642SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_log2(float f)
1643{
1644    return ::log2f(f);
1645}
1646SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_log(float f)
1647{
1648    return ::logf(f);
1649}
1650SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_log10(float f)
1651{
1652    return ::log10f(f);
1653}
1654SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_exp2(float f)
1655{
1656    return ::exp2f(f);
1657}
1658SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_exp(float f)
1659{
1660    return ::expf(f);
1661}
1662SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_abs(float f)
1663{
1664    return ::fabsf(f);
1665}
1666SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_trunc(float f)
1667{
1668    return ::truncf(f);
1669}
1670SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_sqrt(float f)
1671{
1672    return ::sqrtf(f);
1673}
1674SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_rsqrt(float f)
1675{
1676    return ::rsqrtf(f);
1677}
1678SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_sign(float f)
1679{
1680    return (f == 0.0f) ? f : ((f < 0.0f) ? -1.0f : 1.0f);
1681}
1682SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_frac(float f)
1683{
1684    return f - F32_floor(f);
1685}
1686
1687SLANG_FORCE_INLINE SLANG_CUDA_CALL bool F32_isnan(float f)
1688{
1689    return isnan(f);
1690}
1691SLANG_FORCE_INLINE SLANG_CUDA_CALL bool F32_isfinite(float f)
1692{
1693    return isfinite(f);
1694}
1695SLANG_FORCE_INLINE SLANG_CUDA_CALL bool F32_isinf(float f)
1696{
1697    return isinf(f);
1698}
1699
1700// Binary
1701SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_min(float a, float b)
1702{
1703    return ::fminf(a, b);
1704}
1705SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_max(float a, float b)
1706{
1707    return ::fmaxf(a, b);
1708}
1709SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_pow(float a, float b)
1710{
1711    return ::powf(a, b);
1712}
1713SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_fmod(float a, float b)
1714{
1715    return ::fmodf(a, b);
1716}
1717SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_remainder(float a, float b)
1718{
1719    return ::remainderf(a, b);
1720}
1721SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_atan2(float a, float b)
1722{
1723    return float(::atan2(a, b));
1724}
1725
1726SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_frexp(float x, int* e)
1727{
1728    return frexpf(x, e);
1729}
1730
1731SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_modf(float x, float* ip)
1732{
1733    return ::modff(x, ip);
1734}
1735
1736SLANG_FORCE_INLINE SLANG_CUDA_CALL uint32_t F32_asuint(float f)
1737{
1738    Union32 u;
1739    u.f = f;
1740    return u.u;
1741}
1742SLANG_FORCE_INLINE SLANG_CUDA_CALL int32_t F32_asint(float f)
1743{
1744    Union32 u;
1745    u.f = f;
1746    return u.i;
1747}
1748
1749// Ternary
1750SLANG_FORCE_INLINE SLANG_CUDA_CALL float F32_fma(float a, float b, float c)
1751{
1752    return ::fmaf(a, b, c);
1753}
1754
1755
1756// ----------------------------- F64 -----------------------------------------
1757
1758// Unary
1759SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_ceil(double f)
1760{
1761    return ::ceil(f);
1762}
1763SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_floor(double f)
1764{
1765    return ::floor(f);
1766}
1767SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_round(double f)
1768{
1769    return ::round(f);
1770}
1771SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_sin(double f)
1772{
1773    return ::sin(f);
1774}
1775SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_cos(double f)
1776{
1777    return ::cos(f);
1778}
1779SLANG_FORCE_INLINE SLANG_CUDA_CALL void F64_sincos(double f, double* s, double* c)
1780{
1781    ::sincos(f, s, c);
1782}
1783SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_tan(double f)
1784{
1785    return ::tan(f);
1786}
1787SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_asin(double f)
1788{
1789    return ::asin(f);
1790}
1791SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_acos(double f)
1792{
1793    return ::acos(f);
1794}
1795SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_atan(double f)
1796{
1797    return ::atan(f);
1798}
1799SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_sinh(double f)
1800{
1801    return ::sinh(f);
1802}
1803SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_cosh(double f)
1804{
1805    return ::cosh(f);
1806}
1807SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_tanh(double f)
1808{
1809    return ::tanh(f);
1810}
1811SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_log2(double f)
1812{
1813    return ::log2(f);
1814}
1815SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_log(double f)
1816{
1817    return ::log(f);
1818}
1819SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_log10(float f)
1820{
1821    return ::log10(f);
1822}
1823SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_exp2(double f)
1824{
1825    return ::exp2(f);
1826}
1827SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_exp(double f)
1828{
1829    return ::exp(f);
1830}
1831SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_abs(double f)
1832{
1833    return ::fabs(f);
1834}
1835SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_trunc(double f)
1836{
1837    return ::trunc(f);
1838}
1839SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_sqrt(double f)
1840{
1841    return ::sqrt(f);
1842}
1843SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_rsqrt(double f)
1844{
1845    return ::rsqrt(f);
1846}
1847SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_sign(double f)
1848{
1849    return (f == 0.0) ? f : ((f < 0.0) ? -1.0 : 1.0);
1850}
1851SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_frac(double f)
1852{
1853    return f - F64_floor(f);
1854}
1855
1856SLANG_FORCE_INLINE SLANG_CUDA_CALL bool F64_isnan(double f)
1857{
1858    return isnan(f);
1859}
1860SLANG_FORCE_INLINE SLANG_CUDA_CALL bool F64_isfinite(double f)
1861{
1862    return isfinite(f);
1863}
1864SLANG_FORCE_INLINE SLANG_CUDA_CALL bool F64_isinf(double f)
1865{
1866    return isinf(f);
1867}
1868
1869// Binary
1870SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_min(double a, double b)
1871{
1872    return ::fmin(a, b);
1873}
1874SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_max(double a, double b)
1875{
1876    return ::fmax(a, b);
1877}
1878SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_pow(double a, double b)
1879{
1880    return ::pow(a, b);
1881}
1882SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_fmod(double a, double b)
1883{
1884    return ::fmod(a, b);
1885}
1886SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_remainder(double a, double b)
1887{
1888    return ::remainder(a, b);
1889}
1890SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_atan2(double a, double b)
1891{
1892    return ::atan2(a, b);
1893}
1894
1895SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_frexp(double x, int* e)
1896{
1897    return ::frexp(x, e);
1898}
1899
1900SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_modf(double x, double* ip)
1901{
1902    return ::modf(x, ip);
1903}
1904
1905SLANG_FORCE_INLINE SLANG_CUDA_CALL void F64_asuint(double d, uint32_t* low, uint32_t* hi)
1906{
1907    Union64 u;
1908    u.d = d;
1909    *low = uint32_t(u.u);
1910    *hi = uint32_t(u.u >> 32);
1911}
1912
1913SLANG_FORCE_INLINE SLANG_CUDA_CALL void F64_asint(double d, int32_t* low, int32_t* hi)
1914{
1915    Union64 u;
1916    u.d = d;
1917    *low = int32_t(u.u);
1918    *hi = int32_t(u.u >> 32);
1919}
1920
1921// Ternary
1922SLANG_FORCE_INLINE SLANG_CUDA_CALL double F64_fma(double a, double b, double c)
1923{
1924    return ::fma(a, b, c);
1925}
1926
1927// ----------------------------- U8 -----------------------------------------
1928
1929SLANG_FORCE_INLINE SLANG_CUDA_CALL uint32_t U8_countbits(uint8_t v)
1930{
1931    // No native 8bit popc yet, just cast and use 32bit variant
1932    return __popc(uint32_t(v));
1933}
1934
1935// ----------------------------- I8 -----------------------------------------
1936
1937SLANG_FORCE_INLINE SLANG_CUDA_CALL uint32_t I8_countbits(int8_t v)
1938{
1939    return U8_countbits(uint8_t(v));
1940}
1941
1942// ----------------------------- U16 -----------------------------------------
1943
1944SLANG_FORCE_INLINE SLANG_CUDA_CALL uint32_t U16_countbits(uint16_t v)
1945{
1946    // No native 16bit popc yet, just cast and use 32bit variant
1947    return __popc(uint32_t(v));
1948}
1949
1950// ----------------------------- I16 -----------------------------------------
1951
1952SLANG_FORCE_INLINE SLANG_CUDA_CALL uint32_t I16_countbits(int16_t v)
1953{
1954    return U16_countbits(uint16_t(v));
1955}
1956
1957// ----------------------------- U32 -----------------------------------------
1958
1959// Unary
1960SLANG_FORCE_INLINE SLANG_CUDA_CALL uint32_t U32_abs(uint32_t f)
1961{
1962    return f;
1963}
1964
1965// Binary
1966SLANG_FORCE_INLINE SLANG_CUDA_CALL uint32_t U32_min(uint32_t a, uint32_t b)
1967{
1968    return a < b ? a : b;
1969}
1970SLANG_FORCE_INLINE SLANG_CUDA_CALL uint32_t U32_max(uint32_t a, uint32_t b)
1971{
1972    return a > b ? a : b;
1973}
1974
1975SLANG_FORCE_INLINE SLANG_CUDA_CALL float U32_asfloat(uint32_t x)
1976{
1977    Union32 u;
1978    u.u = x;
1979    return u.f;
1980}
1981SLANG_FORCE_INLINE SLANG_CUDA_CALL uint32_t U32_asint(int32_t x)
1982{
1983    return uint32_t(x);
1984}
1985
1986SLANG_FORCE_INLINE SLANG_CUDA_CALL double U32_asdouble(uint32_t low, uint32_t hi)
1987{
1988    Union64 u;
1989    u.u = (uint64_t(hi) << 32) | low;
1990    return u.d;
1991}
1992
1993SLANG_FORCE_INLINE SLANG_CUDA_CALL uint32_t U32_countbits(uint32_t v)
1994{
1995    return __popc(v);
1996}
1997
1998SLANG_FORCE_INLINE SLANG_CUDA_CALL uint32_t U32_firstbitlow(uint32_t v)
1999{
2000    // __ffs returns 1-based bit position or 0 if no bits set
2001    // firstbitlow should return 0-based bit position or ~0u if no bits set
2002    return v == 0 ? ~0u : (__ffs(v) - 1);
2003}
2004
2005SLANG_FORCE_INLINE SLANG_CUDA_CALL uint32_t U32_firstbithigh(uint32_t v)
2006{
2007    // maps to hlsl firstbithigh
2008    if ((int32_t)v < 0)
2009        v = ~v;
2010    if (v == 0)
2011        return ~0u;
2012    return 31 - __clz(v);
2013}
2014
2015// ----------------------------- I32 -----------------------------------------
2016
2017// Unary
2018SLANG_FORCE_INLINE SLANG_CUDA_CALL int32_t I32_abs(int32_t f)
2019{
2020    return (f < 0) ? -f : f;
2021}
2022
2023// Binary
2024SLANG_FORCE_INLINE SLANG_CUDA_CALL int32_t I32_min(int32_t a, int32_t b)
2025{
2026    return a < b ? a : b;
2027}
2028SLANG_FORCE_INLINE SLANG_CUDA_CALL int32_t I32_max(int32_t a, int32_t b)
2029{
2030    return a > b ? a : b;
2031}
2032
2033SLANG_FORCE_INLINE SLANG_CUDA_CALL float I32_asfloat(int32_t x)
2034{
2035    Union32 u;
2036    u.i = x;
2037    return u.f;
2038}
2039SLANG_FORCE_INLINE SLANG_CUDA_CALL uint32_t I32_asuint(int32_t x)
2040{
2041    return uint32_t(x);
2042}
2043SLANG_FORCE_INLINE SLANG_CUDA_CALL double I32_asdouble(int32_t low, int32_t hi)
2044{
2045    Union64 u;
2046    u.u = (uint64_t(hi) << 32) | uint32_t(low);
2047    return u.d;
2048}
2049
2050SLANG_FORCE_INLINE SLANG_CUDA_CALL uint32_t I32_countbits(int32_t v)
2051{
2052    return U32_countbits(uint32_t(v));
2053}
2054
2055SLANG_FORCE_INLINE SLANG_CUDA_CALL uint32_t I32_firstbitlow(int32_t v)
2056{
2057    return U32_firstbitlow(uint32_t(v));
2058}
2059
2060SLANG_FORCE_INLINE SLANG_CUDA_CALL uint32_t I32_firstbithigh(int32_t v)
2061{
2062    return U32_firstbithigh(uint32_t(v));
2063}
2064
2065// ----------------------------- U64 -----------------------------------------
2066
2067SLANG_FORCE_INLINE SLANG_CUDA_CALL int64_t U64_abs(uint64_t f)
2068{
2069    return f;
2070}
2071
2072SLANG_FORCE_INLINE SLANG_CUDA_CALL int64_t U64_min(uint64_t a, uint64_t b)
2073{
2074    return a < b ? a : b;
2075}
2076SLANG_FORCE_INLINE SLANG_CUDA_CALL int64_t U64_max(uint64_t a, uint64_t b)
2077{
2078    return a > b ? a : b;
2079}
2080
2081SLANG_FORCE_INLINE SLANG_CUDA_CALL uint32_t U64_countbits(uint64_t v)
2082{
2083    return __popcll(v);
2084}
2085
2086// ----------------------------- I64 -----------------------------------------
2087
2088SLANG_FORCE_INLINE SLANG_CUDA_CALL int64_t I64_abs(int64_t f)
2089{
2090    return (f < 0) ? -f : f;
2091}
2092
2093SLANG_FORCE_INLINE SLANG_CUDA_CALL int64_t I64_min(int64_t a, int64_t b)
2094{
2095    return a < b ? a : b;
2096}
2097SLANG_FORCE_INLINE SLANG_CUDA_CALL int64_t I64_max(int64_t a, int64_t b)
2098{
2099    return a > b ? a : b;
2100}
2101
2102SLANG_FORCE_INLINE SLANG_CUDA_CALL uint32_t I64_countbits(int64_t v)
2103{
2104    return U64_countbits(uint64_t(v));
2105}
2106
2107// ----------------------------- IPTR -----------------------------------------
2108
2109SLANG_FORCE_INLINE SLANG_CUDA_CALL intptr_t IPTR_abs(intptr_t f)
2110{
2111    return (f < 0) ? -f : f;
2112}
2113
2114SLANG_FORCE_INLINE SLANG_CUDA_CALL intptr_t IPTR_min(intptr_t a, intptr_t b)
2115{
2116    return a < b ? a : b;
2117}
2118
2119SLANG_FORCE_INLINE SLANG_CUDA_CALL intptr_t IPTR_max(intptr_t a, intptr_t b)
2120{
2121    return a > b ? a : b;
2122}
2123
2124// ----------------------------- UPTR -----------------------------------------
2125
2126SLANG_FORCE_INLINE SLANG_CUDA_CALL uintptr_t UPTR_abs(uintptr_t f)
2127{
2128    return f;
2129}
2130
2131SLANG_FORCE_INLINE SLANG_CUDA_CALL uintptr_t UPTR_min(uintptr_t a, uintptr_t b)
2132{
2133    return a < b ? a : b;
2134}
2135
2136SLANG_FORCE_INLINE SLANG_CUDA_CALL uintptr_t UPTR_max(uintptr_t a, uintptr_t b)
2137{
2138    return a > b ? a : b;
2139}
2140
2141// ----------------------------- ResourceType -----------------------------------------
2142
2143
2144// https://docs.microsoft.com/en-us/windows/win32/direct3dhlsl/sm5-object-structuredbuffer-getdimensions
2145// Missing  Load(_In_  int  Location, _Out_ uint Status);
2146
2147template<typename T>
2148struct StructuredBuffer
2149{
2150    SLANG_CUDA_CALL T& operator[](size_t index) const
2151    {
2152#ifndef SLANG_CUDA_STRUCTURED_BUFFER_NO_COUNT
2153        SLANG_BOUND_CHECK(index, count);
2154#endif
2155        return data[index];
2156    }
2157
2158    SLANG_CUDA_CALL T& Load(size_t index) const
2159    {
2160#ifndef SLANG_CUDA_STRUCTURED_BUFFER_NO_COUNT
2161        SLANG_BOUND_CHECK(index, count);
2162#endif
2163        return data[index];
2164    }
2165
2166#ifndef SLANG_CUDA_STRUCTURED_BUFFER_NO_COUNT
2167    SLANG_CUDA_CALL void GetDimensions(uint32_t* outNumStructs, uint32_t* outStride) const
2168    {
2169        *outNumStructs = uint32_t(count);
2170        *outStride = uint32_t(sizeof(T));
2171    }
2172#endif
2173
2174    T* data;
2175#ifndef SLANG_CUDA_STRUCTURED_BUFFER_NO_COUNT
2176    size_t count;
2177#endif
2178};
2179
2180template<typename T>
2181struct RWStructuredBuffer : StructuredBuffer<T>
2182{
2183    SLANG_CUDA_CALL T& operator[](size_t index) const
2184    {
2185#ifndef SLANG_CUDA_STRUCTURED_BUFFER_NO_COUNT
2186        SLANG_BOUND_CHECK(index, this->count);
2187#endif
2188        return this->data[index];
2189    }
2190};
2191
2192// Missing  Load(_In_  int  Location, _Out_ uint Status);
2193struct ByteAddressBuffer
2194{
2195    SLANG_CUDA_CALL void GetDimensions(uint32_t* outDim) const { *outDim = uint32_t(sizeInBytes); }
2196    SLANG_CUDA_CALL uint32_t Load(size_t index) const
2197    {
2198        SLANG_BOUND_CHECK_BYTE_ADDRESS(index, 4, sizeInBytes);
2199        return data[index >> 2];
2200    }
2201    SLANG_CUDA_CALL uint2 Load2(size_t index) const
2202    {
2203        SLANG_BOUND_CHECK_BYTE_ADDRESS(index, 8, sizeInBytes);
2204        const size_t dataIdx = index >> 2;
2205        return uint2{data[dataIdx], data[dataIdx + 1]};
2206    }
2207    SLANG_CUDA_CALL uint3 Load3(size_t index) const
2208    {
2209        SLANG_BOUND_CHECK_BYTE_ADDRESS(index, 12, sizeInBytes);
2210        const size_t dataIdx = index >> 2;
2211        return uint3{data[dataIdx], data[dataIdx + 1], data[dataIdx + 2]};
2212    }
2213    SLANG_CUDA_CALL uint4 Load4(size_t index) const
2214    {
2215        SLANG_BOUND_CHECK_BYTE_ADDRESS(index, 16, sizeInBytes);
2216        const size_t dataIdx = index >> 2;
2217        return uint4{data[dataIdx], data[dataIdx + 1], data[dataIdx + 2], data[dataIdx + 3]};
2218    }
2219    template<typename T>
2220    SLANG_CUDA_CALL T Load(size_t index) const
2221    {
2222        SLANG_BOUND_CHECK_BYTE_ADDRESS(index, sizeof(T), sizeInBytes);
2223        T data;
2224        memcpy(&data, ((const char*)this->data) + index, sizeof(T));
2225        return data;
2226    }
2227    template<typename T>
2228    SLANG_CUDA_CALL StructuredBuffer<T> asStructuredBuffer() const
2229    {
2230        StructuredBuffer<T> rs;
2231        rs.data = (T*)data;
2232        rs.count = sizeInBytes / sizeof(T);
2233        return rs;
2234    }
2235    const uint32_t* data;
2236    size_t sizeInBytes; //< Must be multiple of 4
2237};
2238
2239// https://docs.microsoft.com/en-us/windows/win32/direct3dhlsl/sm5-object-rwbyteaddressbuffer
2240// Atomic operations support
2241
2242// Signed 64-bit atomic wrappers
2243// CUDA only supports unsigned long long atomics, so we cast signed to unsigned
2244__device__ __forceinline__ long long atomicExch(long long* address, long long val)
2245{
2246    return (long long)atomicExch((unsigned long long*)address, (unsigned long long)val);
2247}
2248
2249__device__ __forceinline__ long long atomicCAS(long long* address, long long compare, long long val)
2250{
2251    return (long long)atomicCAS(
2252        (unsigned long long*)address,
2253        (unsigned long long)compare,
2254        (unsigned long long)val);
2255}
2256
2257// Float bitwise atomic compare-and-swap
2258// Uses integer atomics to preserve exact float bit patterns
2259__device__ __forceinline__ float atomicCAS(float* address, float compare, float val)
2260{
2261    int* addr_as_int = (int*)address;
2262    int old = atomicCAS(addr_as_int, __float_as_int(compare), __float_as_int(val));
2263    return __int_as_float(old);
2264}
2265
2266// Missing support for Load with status
2267struct RWByteAddressBuffer
2268{
2269    SLANG_CUDA_CALL void GetDimensions(uint32_t* outDim) const { *outDim = uint32_t(sizeInBytes); }
2270
2271    SLANG_CUDA_CALL uint32_t Load(size_t index) const
2272    {
2273        SLANG_BOUND_CHECK_BYTE_ADDRESS(index, 4, sizeInBytes);
2274        return data[index >> 2];
2275    }
2276    SLANG_CUDA_CALL uint2 Load2(size_t index) const
2277    {
2278        SLANG_BOUND_CHECK_BYTE_ADDRESS(index, 8, sizeInBytes);
2279        const size_t dataIdx = index >> 2;
2280        return uint2{data[dataIdx], data[dataIdx + 1]};
2281    }
2282    SLANG_CUDA_CALL uint3 Load3(size_t index) const
2283    {
2284        SLANG_BOUND_CHECK_BYTE_ADDRESS(index, 12, sizeInBytes);
2285        const size_t dataIdx = index >> 2;
2286        return uint3{data[dataIdx], data[dataIdx + 1], data[dataIdx + 2]};
2287    }
2288    SLANG_CUDA_CALL uint4 Load4(size_t index) const
2289    {
2290        SLANG_BOUND_CHECK_BYTE_ADDRESS(index, 16, sizeInBytes);
2291        const size_t dataIdx = index >> 2;
2292        return uint4{data[dataIdx], data[dataIdx + 1], data[dataIdx + 2], data[dataIdx + 3]};
2293    }
2294    template<typename T>
2295    SLANG_CUDA_CALL T Load(size_t index) const
2296    {
2297        SLANG_BOUND_CHECK_BYTE_ADDRESS(index, sizeof(T), sizeInBytes);
2298        T data;
2299        memcpy(&data, ((const char*)this->data) + index, sizeof(T));
2300        return data;
2301    }
2302
2303    SLANG_CUDA_CALL void Store(size_t index, uint32_t v) const
2304    {
2305        SLANG_BOUND_CHECK_BYTE_ADDRESS(index, 4, sizeInBytes);
2306        data[index >> 2] = v;
2307    }
2308    SLANG_CUDA_CALL void Store2(size_t index, uint2 v) const
2309    {
2310        SLANG_BOUND_CHECK_BYTE_ADDRESS(index, 8, sizeInBytes);
2311        const size_t dataIdx = index >> 2;
2312        data[dataIdx + 0] = v.x;
2313        data[dataIdx + 1] = v.y;
2314    }
2315    SLANG_CUDA_CALL void Store3(size_t index, uint3 v) const
2316    {
2317        SLANG_BOUND_CHECK_BYTE_ADDRESS(index, 12, sizeInBytes);
2318        const size_t dataIdx = index >> 2;
2319        data[dataIdx + 0] = v.x;
2320        data[dataIdx + 1] = v.y;
2321        data[dataIdx + 2] = v.z;
2322    }
2323    SLANG_CUDA_CALL void Store4(size_t index, uint4 v) const
2324    {
2325        SLANG_BOUND_CHECK_BYTE_ADDRESS(index, 16, sizeInBytes);
2326        const size_t dataIdx = index >> 2;
2327        data[dataIdx + 0] = v.x;
2328        data[dataIdx + 1] = v.y;
2329        data[dataIdx + 2] = v.z;
2330        data[dataIdx + 3] = v.w;
2331    }
2332    template<typename T>
2333    SLANG_CUDA_CALL void Store(size_t index, T const& value) const
2334    {
2335        SLANG_BOUND_CHECK_BYTE_ADDRESS(index, sizeof(T), sizeInBytes);
2336        memcpy((char*)data + index, &value, sizeof(T));
2337    }
2338
2339    /// Can be used in the core module to gain access
2340    template<typename T>
2341    SLANG_CUDA_CALL T* _getPtrAt(size_t index)
2342    {
2343        SLANG_BOUND_CHECK_BYTE_ADDRESS(index, sizeof(T), sizeInBytes);
2344        return (T*)(((char*)data) + index);
2345    }
2346    template<typename T>
2347    SLANG_CUDA_CALL RWStructuredBuffer<T> asStructuredBuffer() const
2348    {
2349        RWStructuredBuffer<T> rs;
2350        rs.data = (T*)data;
2351        rs.count = sizeInBytes / sizeof(T);
2352        return rs;
2353    }
2354    uint32_t* data;
2355    size_t sizeInBytes; //< Must be multiple of 4
2356};
2357
2358
2359// ---------------------- Wave --------------------------------------
2360
2361// TODO(JS): It appears that cuda does not have a simple way to get a lane index.
2362//
2363// Another approach could be...
2364// laneId = ((threadIdx.z * blockDim.y + threadIdx.y) * blockDim.x + threadIdx.x) &
2365// SLANG_CUDA_WARP_MASK If that is really true another way to do this, would be for code generator
2366// to add this function with the [numthreads] baked in.
2367//
2368// For now I'll just assume you have a launch that makes the following correct if the kernel uses
2369// WaveGetLaneIndex()
2370#ifndef SLANG_USE_ASM_LANE_ID
2371__forceinline__ __device__ uint32_t _getLaneId()
2372{
2373    // If the launch is (or I guess some multiple of the warp size)
2374    // we try this mechanism, which is apparently faster.
2375    return threadIdx.x & SLANG_CUDA_WARP_MASK;
2376}
2377#else
2378__forceinline__ __device__ uint32_t _getLaneId()
2379{
2380    // https://stackoverflow.com/questions/44337309/whats-the-most-efficient-way-to-calculate-the-warp-id-lane-id-in-a-1-d-grid#
2381    // This mechanism is not the fastest way to do it, and that is why the other mechanism
2382    // is the default. But the other mechanism relies on a launch that makes the assumption
2383    // true.
2384    unsigned ret;
2385    asm volatile("mov.u32 %0, %laneid;" : "=r"(ret));
2386    return ret;
2387}
2388#endif
2389
2390typedef int WarpMask;
2391
2392// It appears that the __activemask() cannot always be used because
2393// threads need to be converged.
2394//
2395// For CUDA the article claims mask has to be used carefully
2396// https://devblogs.nvidia.com/using-cuda-warp-level-primitives/
2397// With the Warp intrinsics there is no mask, and it's just the 'active lanes'.
2398// __activemask() though does not require there is convergence, so that doesn't work.
2399//
2400// '__ballot_sync' produces a convergance.
2401//
2402// From the CUDA docs:
2403// ```For __all_sync, __any_sync, and __ballot_sync, a mask must be passed that specifies the
2404// threads participating in the call. A bit, representing the thread's lane ID, must be set for each
2405// participating thread to ensure they are properly converged before the intrinsic is executed by
2406// the hardware. All active threads named in mask must execute the same intrinsic with the same
2407// mask, or the result is undefined.```
2408//
2409// Currently there isn't a mechanism to correctly get the mask without it being passed through.
2410// Doing so will most likely require some changes to slang code generation to track masks, for now
2411// then we use _getActiveMask.
2412
2413// Return mask of all the lanes less than the current lane
2414__forceinline__ __device__ WarpMask _getLaneLtMask()
2415{
2416    return (int(1) << _getLaneId()) - 1;
2417}
2418
2419// TODO(JS):
2420// THIS IS NOT CORRECT! That determining the appropriate active mask requires appropriate
2421// mask tracking.
2422__forceinline__ __device__ WarpMask _getActiveMask()
2423{
2424    return __ballot_sync(__activemask(), true);
2425}
2426
2427// Return a mask suitable for the 'MultiPrefix' style functions
2428__forceinline__ __device__ WarpMask _getMultiPrefixMask(int mask)
2429{
2430    return mask;
2431}
2432
2433// Note! Note will return true if mask is 0, but thats okay, because there must be one
2434// lane active to execute anything
2435__inline__ __device__ bool _waveIsSingleLane(WarpMask mask)
2436{
2437    return (mask & (mask - 1)) == 0;
2438}
2439
2440// Returns the power of 2 size of run of set bits. Returns 0 if not a suitable run.
2441// Examples:
2442// 0b00000000'00000000'00000000'11111111 -> 8
2443// 0b11111111'11111111'11111111'11111111 -> 32
2444// 0b00000000'00000000'00000000'00011111 -> 0 (since 5 is not a power of 2)
2445// 0b00000000'00000000'00000000'11110000 -> 0 (since the run of bits does not start at the LSB)
2446// 0b00000000'00000000'00000000'00100111 -> 0 (since it is not a single contiguous run)
2447__inline__ __device__ int _waveCalcPow2Offset(WarpMask mask)
2448{
2449    // This should be the most common case, so fast path it
2450    if (mask == SLANG_CUDA_WARP_BITMASK)
2451    {
2452        return SLANG_CUDA_WARP_SIZE;
2453    }
2454    // Is it a contiguous run of bits?
2455    if ((mask & (mask + 1)) == 0)
2456    {
2457        // const int offsetSize = __ffs(mask + 1) - 1;
2458        const int offset = 32 - __clz(mask);
2459        // Is it a power of 2 size
2460        if ((offset & (offset - 1)) == 0)
2461        {
2462            return offset;
2463        }
2464    }
2465    return 0;
2466}
2467
2468__inline__ __device__ bool _waveIsFirstLane()
2469{
2470    const WarpMask mask = __activemask();
2471    // We special case bit 0, as that most warps are expected to be fully active.
2472
2473    // mask & -mask, isolates the lowest set bit.
2474    // return (mask & 1 ) || ((mask & -mask) == (1 << _getLaneId()));
2475
2476    // This mechanism is most similar to what was in an nVidia post, so assume it is prefered.
2477    return (mask & 1) || ((__ffs(mask) - 1) == _getLaneId());
2478}
2479
2480template<typename T>
2481struct WaveOpOr
2482{
2483    __inline__ __device__ static T getInitial(T a) { return 0; }
2484    __inline__ __device__ static T doOp(T a, T b) { return a | b; }
2485};
2486
2487template<typename T>
2488struct WaveOpAnd
2489{
2490    __inline__ __device__ static T getInitial(T a) { return ~T(0); }
2491    __inline__ __device__ static T doOp(T a, T b) { return a & b; }
2492};
2493
2494template<typename T>
2495struct WaveOpXor
2496{
2497    __inline__ __device__ static T getInitial(T a) { return 0; }
2498    __inline__ __device__ static T doOp(T a, T b) { return a ^ b; }
2499    __inline__ __device__ static T doInverse(T a, T b) { return a ^ b; }
2500};
2501
2502template<typename T>
2503struct WaveOpAdd
2504{
2505    __inline__ __device__ static T getInitial(T a) { return 0; }
2506    __inline__ __device__ static T doOp(T a, T b) { return a + b; }
2507    __inline__ __device__ static T doInverse(T a, T b) { return a - b; }
2508};
2509
2510template<typename T>
2511struct WaveOpMul
2512{
2513    __inline__ __device__ static T getInitial(T a) { return T(1); }
2514    __inline__ __device__ static T doOp(T a, T b) { return a * b; }
2515    // Using this inverse for int is probably undesirable - because in general it requires T to have
2516    // more precision There is also a performance aspect to it, where divides are generally
2517    // significantly slower
2518    __inline__ __device__ static T doInverse(T a, T b) { return a / b; }
2519};
2520
2521template<typename T>
2522struct WaveOpMax
2523{
2524    __inline__ __device__ static T getInitial(T a, bool exclusive = false);
2525    __inline__ __device__ static T doOp(T a, T b) { return a > b ? a : b; }
2526};
2527
2528template<typename T>
2529struct WaveOpMin
2530{
2531    __inline__ __device__ static T getInitial(T a, bool exclusive = false);
2532    __inline__ __device__ static T doOp(T a, T b) { return a < b ? a : b; }
2533};
2534
2535// Compact specializations using macro for getInitial
2536#define SLANG_WAVE_MIN_SPEC(T, EXCL_VAL)                                  \
2537    template<>                                                            \
2538    __inline__ __device__ T WaveOpMin<T>::getInitial(T a, bool exclusive) \
2539    {                                                                     \
2540        return exclusive ? (EXCL_VAL) : a;                                \
2541    }
2542
2543#define SLANG_WAVE_MAX_SPEC(T, EXCL_VAL)                                  \
2544    template<>                                                            \
2545    __inline__ __device__ T WaveOpMax<T>::getInitial(T a, bool exclusive) \
2546    {                                                                     \
2547        return exclusive ? (EXCL_VAL) : a;                                \
2548    }
2549
2550// Min specializations (exclusive identity = max value)
2551SLANG_WAVE_MIN_SPEC(float, SLANG_INFINITY)
2552SLANG_WAVE_MIN_SPEC(double, SLANG_INFINITY)
2553SLANG_WAVE_MIN_SPEC(int, 0x7FFFFFFF)
2554SLANG_WAVE_MIN_SPEC(uint, 0xFFFFFFFF)
2555SLANG_WAVE_MIN_SPEC(char, (char)0x7F)
2556SLANG_WAVE_MIN_SPEC(int8_t, (int8_t)0x7F)
2557SLANG_WAVE_MIN_SPEC(uint8_t, (uint8_t)0xFF)
2558SLANG_WAVE_MIN_SPEC(int16_t, (int16_t)0x7FFF)
2559SLANG_WAVE_MIN_SPEC(uint16_t, (uint16_t)0xFFFF)
2560SLANG_WAVE_MIN_SPEC(int64_t, 0x7FFFFFFFFFFFFFFFLL)
2561SLANG_WAVE_MIN_SPEC(uint64_t, 0xFFFFFFFFFFFFFFFFULL)
2562#if SLANG_CUDA_ENABLE_HALF
2563SLANG_WAVE_MIN_SPEC(__half, __ushort_as_half(0x7BFF))
2564#endif
2565
2566// Max specializations (exclusive identity = min value)
2567SLANG_WAVE_MAX_SPEC(float, -SLANG_INFINITY)
2568SLANG_WAVE_MAX_SPEC(double, -SLANG_INFINITY)
2569SLANG_WAVE_MAX_SPEC(int, (int)0x80000000)
2570SLANG_WAVE_MAX_SPEC(uint, 0)
2571SLANG_WAVE_MAX_SPEC(char, (char)0x80)
2572SLANG_WAVE_MAX_SPEC(int8_t, (int8_t)0x80)
2573SLANG_WAVE_MAX_SPEC(uint8_t, 0)
2574SLANG_WAVE_MAX_SPEC(int16_t, (int16_t)0x8000)
2575SLANG_WAVE_MAX_SPEC(uint16_t, 0)
2576SLANG_WAVE_MAX_SPEC(int64_t, (int64_t)0x8000000000000000LL)
2577SLANG_WAVE_MAX_SPEC(uint64_t, 0)
2578#if SLANG_CUDA_ENABLE_HALF
2579SLANG_WAVE_MAX_SPEC(__half, __ushort_as_half(0xFBFF))
2580#endif
2581
2582#undef SLANG_WAVE_MIN_SPEC
2583#undef SLANG_WAVE_MAX_SPEC
2584
2585template<typename T>
2586struct ElementTypeTrait;
2587
2588// Scalar
2589template<>
2590struct ElementTypeTrait<int>
2591{
2592    typedef int Type;
2593};
2594template<>
2595struct ElementTypeTrait<uint>
2596{
2597    typedef uint Type;
2598};
2599template<>
2600struct ElementTypeTrait<float>
2601{
2602    typedef float Type;
2603};
2604template<>
2605struct ElementTypeTrait<double>
2606{
2607    typedef double Type;
2608};
2609template<>
2610struct ElementTypeTrait<uint64_t>
2611{
2612    typedef uint64_t Type;
2613};
2614template<>
2615struct ElementTypeTrait<int64_t>
2616{
2617    typedef int64_t Type;
2618};
2619template<>
2620struct ElementTypeTrait<char>
2621{
2622    typedef char Type;
2623};
2624template<>
2625struct ElementTypeTrait<uchar>
2626{
2627    typedef uchar Type;
2628};
2629template<>
2630struct ElementTypeTrait<short>
2631{
2632    typedef short Type;
2633};
2634template<>
2635struct ElementTypeTrait<ushort>
2636{
2637    typedef ushort Type;
2638};
2639#if SLANG_CUDA_ENABLE_HALF
2640template<>
2641struct ElementTypeTrait<__half>
2642{
2643    typedef __half Type;
2644};
2645#endif
2646
2647// Vector
2648template<>
2649struct ElementTypeTrait<int1>
2650{
2651    typedef int Type;
2652};
2653template<>
2654struct ElementTypeTrait<int2>
2655{
2656    typedef int Type;
2657};
2658template<>
2659struct ElementTypeTrait<int3>
2660{
2661    typedef int Type;
2662};
2663template<>
2664struct ElementTypeTrait<int4>
2665{
2666    typedef int Type;
2667};
2668
2669template<>
2670struct ElementTypeTrait<uint1>
2671{
2672    typedef uint Type;
2673};
2674template<>
2675struct ElementTypeTrait<uint2>
2676{
2677    typedef uint Type;
2678};
2679template<>
2680struct ElementTypeTrait<uint3>
2681{
2682    typedef uint Type;
2683};
2684template<>
2685struct ElementTypeTrait<uint4>
2686{
2687    typedef uint Type;
2688};
2689
2690template<>
2691struct ElementTypeTrait<float1>
2692{
2693    typedef float Type;
2694};
2695template<>
2696struct ElementTypeTrait<float2>
2697{
2698    typedef float Type;
2699};
2700template<>
2701struct ElementTypeTrait<float3>
2702{
2703    typedef float Type;
2704};
2705template<>
2706struct ElementTypeTrait<float4>
2707{
2708    typedef float Type;
2709};
2710
2711template<>
2712struct ElementTypeTrait<double1>
2713{
2714    typedef double Type;
2715};
2716template<>
2717struct ElementTypeTrait<double2>
2718{
2719    typedef double Type;
2720};
2721template<>
2722struct ElementTypeTrait<double3>
2723{
2724    typedef double Type;
2725};
2726template<>
2727struct ElementTypeTrait<double4>
2728{
2729    typedef double Type;
2730};
2731
2732// Additional vector types
2733template<>
2734struct ElementTypeTrait<char2>
2735{
2736    typedef char Type;
2737};
2738template<>
2739struct ElementTypeTrait<char3>
2740{
2741    typedef char Type;
2742};
2743template<>
2744struct ElementTypeTrait<char4>
2745{
2746    typedef char Type;
2747};
2748template<>
2749struct ElementTypeTrait<uchar2>
2750{
2751    typedef uchar Type;
2752};
2753template<>
2754struct ElementTypeTrait<uchar3>
2755{
2756    typedef uchar Type;
2757};
2758template<>
2759struct ElementTypeTrait<uchar4>
2760{
2761    typedef uchar Type;
2762};
2763template<>
2764struct ElementTypeTrait<short2>
2765{
2766    typedef short Type;
2767};
2768template<>
2769struct ElementTypeTrait<short3>
2770{
2771    typedef short Type;
2772};
2773template<>
2774struct ElementTypeTrait<short4>
2775{
2776    typedef short Type;
2777};
2778template<>
2779struct ElementTypeTrait<ushort2>
2780{
2781    typedef ushort Type;
2782};
2783template<>
2784struct ElementTypeTrait<ushort3>
2785{
2786    typedef ushort Type;
2787};
2788template<>
2789struct ElementTypeTrait<ushort4>
2790{
2791    typedef ushort Type;
2792};
2793template<>
2794struct ElementTypeTrait<longlong2>
2795{
2796    typedef int64_t Type;
2797};
2798template<>
2799struct ElementTypeTrait<longlong3>
2800{
2801    typedef int64_t Type;
2802};
2803template<>
2804struct ElementTypeTrait<longlong4>
2805{
2806    typedef int64_t Type;
2807};
2808template<>
2809struct ElementTypeTrait<ulonglong2>
2810{
2811    typedef uint64_t Type;
2812};
2813template<>
2814struct ElementTypeTrait<ulonglong3>
2815{
2816    typedef uint64_t Type;
2817};
2818template<>
2819struct ElementTypeTrait<ulonglong4>
2820{
2821    typedef uint64_t Type;
2822};
2823#if SLANG_CUDA_ENABLE_HALF
2824template<>
2825struct ElementTypeTrait<__half2>
2826{
2827    typedef __half Type;
2828};
2829template<>
2830struct ElementTypeTrait<__half3>
2831{
2832    typedef __half Type;
2833};
2834template<>
2835struct ElementTypeTrait<__half4>
2836{
2837    typedef __half Type;
2838};
2839#endif
2840
2841// Matrix
2842template<typename T, int ROWS, int COLS>
2843struct ElementTypeTrait<Matrix<T, ROWS, COLS>>
2844{
2845    typedef T Type;
2846};
2847
2848// Scalar
2849template<typename INTF, typename T>
2850__device__ T _waveReduceScalar(WarpMask mask, T val)
2851{
2852    const int offsetSize = _waveCalcPow2Offset(mask);
2853    if (offsetSize > 0)
2854    {
2855        // Fast path O(log2(activeLanes))
2856        for (int offset = offsetSize >> 1; offset > 0; offset >>= 1)
2857        {
2858            val = INTF::doOp(val, __shfl_xor_sync(mask, val, offset));
2859        }
2860    }
2861    else if (!_waveIsSingleLane(mask))
2862    {
2863        T result = INTF::getInitial(val);
2864        int remaining = mask;
2865        while (remaining)
2866        {
2867            const int laneBit = remaining & -remaining;
2868            // Get the sourceLane
2869            const int srcLane = __ffs(laneBit) - 1;
2870            // Broadcast (can also broadcast to self)
2871            result = INTF::doOp(result, __shfl_sync(mask, val, srcLane));
2872            remaining &= ~laneBit;
2873        }
2874        return result;
2875    }
2876    return val;
2877}
2878
2879
2880// Multiple values
2881template<typename INTF, typename T, size_t COUNT>
2882__device__ void _waveReduceMultiple(WarpMask mask, T* val)
2883{
2884    const int offsetSize = _waveCalcPow2Offset(mask);
2885    if (offsetSize > 0)
2886    {
2887        // Fast path O(log2(activeLanes))
2888        for (int offset = offsetSize >> 1; offset > 0; offset >>= 1)
2889        {
2890            for (size_t i = 0; i < COUNT; ++i)
2891            {
2892                val[i] = INTF::doOp(val[i], __shfl_xor_sync(mask, val[i], offset));
2893            }
2894        }
2895    }
2896    else if (!_waveIsSingleLane(mask))
2897    {
2898        // Copy the original
2899        T originalVal[COUNT];
2900        for (size_t i = 0; i < COUNT; ++i)
2901        {
2902            const T v = val[i];
2903            originalVal[i] = v;
2904            val[i] = INTF::getInitial(v);
2905        }
2906
2907        int remaining = mask;
2908        while (remaining)
2909        {
2910            const int laneBit = remaining & -remaining;
2911            // Get the sourceLane
2912            const int srcLane = __ffs(laneBit) - 1;
2913            // Broadcast (can also broadcast to self)
2914            for (size_t i = 0; i < COUNT; ++i)
2915            {
2916                val[i] = INTF::doOp(val[i], __shfl_sync(mask, originalVal[i], srcLane));
2917            }
2918            remaining &= ~laneBit;
2919        }
2920    }
2921}
2922
2923template<typename INTF, typename T>
2924__device__ void _waveReduceMultiple(WarpMask mask, T* val)
2925{
2926    typedef typename ElementTypeTrait<T>::Type ElemType;
2927    _waveReduceMultiple<INTF, ElemType, sizeof(T) / sizeof(ElemType)>(mask, (ElemType*)val);
2928}
2929
2930template<typename T>
2931__inline__ __device__ T _waveOr(WarpMask mask, T val)
2932{
2933    return _waveReduceScalar<WaveOpOr<T>, T>(mask, val);
2934}
2935
2936template<typename T>
2937__inline__ __device__ T _waveAnd(WarpMask mask, T val)
2938{
2939    return _waveReduceScalar<WaveOpAnd<T>, T>(mask, val);
2940}
2941
2942template<typename T>
2943__inline__ __device__ T _waveXor(WarpMask mask, T val)
2944{
2945    return _waveReduceScalar<WaveOpXor<T>, T>(mask, val);
2946}
2947
2948template<typename T>
2949__inline__ __device__ T _waveProduct(WarpMask mask, T val)
2950{
2951    return _waveReduceScalar<WaveOpMul<T>, T>(mask, val);
2952}
2953
2954template<typename T>
2955__inline__ __device__ T _waveSum(WarpMask mask, T val)
2956{
2957    return _waveReduceScalar<WaveOpAdd<T>, T>(mask, val);
2958}
2959
2960template<typename T>
2961__inline__ __device__ T _waveMin(WarpMask mask, T val)
2962{
2963    return _waveReduceScalar<WaveOpMin<T>, T>(mask, val);
2964}
2965
2966template<typename T>
2967__inline__ __device__ T _waveMax(WarpMask mask, T val)
2968{
2969    return _waveReduceScalar<WaveOpMax<T>, T>(mask, val);
2970}
2971
2972// Fast-path specializations when CUDA warp reduce operators are available
2973#if __CUDA_ARCH__ >= 800 // 8.x or higher
2974template<>
2975__inline__ __device__ unsigned _waveOr<unsigned>(WarpMask mask, unsigned val)
2976{
2977    return __reduce_or_sync(mask, val);
2978}
2979
2980template<>
2981__inline__ __device__ unsigned _waveAnd<unsigned>(WarpMask mask, unsigned val)
2982{
2983    return __reduce_and_sync(mask, val);
2984}
2985
2986template<>
2987__inline__ __device__ unsigned _waveXor<unsigned>(WarpMask mask, unsigned val)
2988{
2989    return __reduce_xor_sync(mask, val);
2990}
2991
2992template<>
2993__inline__ __device__ unsigned _waveSum<unsigned>(WarpMask mask, unsigned val)
2994{
2995    return __reduce_add_sync(mask, val);
2996}
2997
2998template<>
2999__inline__ __device__ int _waveSum<int>(WarpMask mask, int val)
3000{
3001    return __reduce_add_sync(mask, val);
3002}
3003
3004template<>
3005__inline__ __device__ unsigned _waveMin<unsigned>(WarpMask mask, unsigned val)
3006{
3007    return __reduce_min_sync(mask, val);
3008}
3009
3010template<>
3011__inline__ __device__ int _waveMin<int>(WarpMask mask, int val)
3012{
3013    return __reduce_min_sync(mask, val);
3014}
3015
3016template<>
3017__inline__ __device__ unsigned _waveMax<unsigned>(WarpMask mask, unsigned val)
3018{
3019    return __reduce_max_sync(mask, val);
3020}
3021
3022template<>
3023__inline__ __device__ int _waveMax<int>(WarpMask mask, int val)
3024{
3025    return __reduce_max_sync(mask, val);
3026}
3027#endif
3028
3029
3030// Multiple
3031
3032template<typename T>
3033__inline__ __device__ T _waveOrMultiple(WarpMask mask, T val)
3034{
3035    typedef typename ElementTypeTrait<T>::Type ElemType;
3036    _waveReduceMultiple<WaveOpOr<ElemType>>(mask, &val);
3037    return val;
3038}
3039
3040template<typename T>
3041__inline__ __device__ T _waveAndMultiple(WarpMask mask, T val)
3042{
3043    typedef typename ElementTypeTrait<T>::Type ElemType;
3044    _waveReduceMultiple<WaveOpAnd<ElemType>>(mask, &val);
3045    return val;
3046}
3047
3048template<typename T>
3049__inline__ __device__ T _waveXorMultiple(WarpMask mask, T val)
3050{
3051    typedef typename ElementTypeTrait<T>::Type ElemType;
3052    _waveReduceMultiple<WaveOpXor<ElemType>>(mask, &val);
3053    return val;
3054}
3055
3056template<typename T>
3057__inline__ __device__ T _waveProductMultiple(WarpMask mask, T val)
3058{
3059    typedef typename ElementTypeTrait<T>::Type ElemType;
3060    _waveReduceMultiple<WaveOpMul<ElemType>>(mask, &val);
3061    return val;
3062}
3063
3064template<typename T>
3065__inline__ __device__ T _waveSumMultiple(WarpMask mask, T val)
3066{
3067    typedef typename ElementTypeTrait<T>::Type ElemType;
3068    _waveReduceMultiple<WaveOpAdd<ElemType>>(mask, &val);
3069    return val;
3070}
3071
3072template<typename T>
3073__inline__ __device__ T _waveMinMultiple(WarpMask mask, T val)
3074{
3075    typedef typename ElementTypeTrait<T>::Type ElemType;
3076    _waveReduceMultiple<WaveOpMin<ElemType>>(mask, &val);
3077    return val;
3078}
3079
3080template<typename T>
3081__inline__ __device__ T _waveMaxMultiple(WarpMask mask, T val)
3082{
3083    typedef typename ElementTypeTrait<T>::Type ElemType;
3084    _waveReduceMultiple<WaveOpMax<ElemType>>(mask, &val);
3085    return val;
3086}
3087
3088
3089template<typename T>
3090__inline__ __device__ bool _waveAllEqual(WarpMask mask, T val)
3091{
3092    int pred;
3093    __match_all_sync(mask, val, &pred);
3094    return pred != 0;
3095}
3096
3097template<typename T>
3098__inline__ __device__ bool _waveAllEqualMultiple(WarpMask mask, T inVal)
3099{
3100    typedef typename ElementTypeTrait<T>::Type ElemType;
3101    const size_t count = sizeof(T) / sizeof(ElemType);
3102    int pred;
3103    const ElemType* src = (const ElemType*)&inVal;
3104    for (size_t i = 0; i < count; ++i)
3105    {
3106        __match_all_sync(mask, src[i], &pred);
3107        if (pred == 0)
3108        {
3109            return false;
3110        }
3111    }
3112    return true;
3113}
3114
3115template<typename T>
3116__inline__ __device__ T _waveReadFirst(WarpMask mask, T val)
3117{
3118    const int lowestLaneId = __ffs(mask) - 1;
3119    return __shfl_sync(mask, val, lowestLaneId);
3120}
3121
3122template<typename T>
3123__inline__ __device__ T _waveReadFirstMultiple(WarpMask mask, T inVal)
3124{
3125    typedef typename ElementTypeTrait<T>::Type ElemType;
3126    const size_t count = sizeof(T) / sizeof(ElemType);
3127    T outVal;
3128    const ElemType* src = (const ElemType*)&inVal;
3129    ElemType* dst = (ElemType*)&outVal;
3130    const int lowestLaneId = __ffs(mask) - 1;
3131    for (size_t i = 0; i < count; ++i)
3132    {
3133        dst[i] = __shfl_sync(mask, src[i], lowestLaneId);
3134    }
3135    return outVal;
3136}
3137
3138template<typename T>
3139__inline__ __device__ T _waveShuffleMultiple(WarpMask mask, T inVal, int lane)
3140{
3141    typedef typename ElementTypeTrait<T>::Type ElemType;
3142    const size_t count = sizeof(T) / sizeof(ElemType);
3143    T outVal;
3144    const ElemType* src = (const ElemType*)&inVal;
3145    ElemType* dst = (ElemType*)&outVal;
3146    for (size_t i = 0; i < count; ++i)
3147    {
3148        dst[i] = __shfl_sync(mask, src[i], lane);
3149    }
3150    return outVal;
3151}
3152
3153// Scalar
3154
3155// Invertable means that when we get to the end of the reduce, we can remove val (to make
3156// exclusive), using the inverse of the op.
3157template<typename INTF, typename T>
3158__device__ T _wavePrefixInvertableScalar(WarpMask mask, T val)
3159{
3160    const int offsetSize = _waveCalcPow2Offset(mask);
3161
3162    const int laneId = _getLaneId();
3163    T result;
3164    if (offsetSize > 0)
3165    {
3166        // Sum is calculated inclusive of this lanes value
3167        result = val;
3168        for (int i = 1; i < offsetSize; i += i)
3169        {
3170            const T readVal = __shfl_up_sync(mask, result, i, offsetSize);
3171            if (laneId >= i)
3172            {
3173                result = INTF::doOp(result, readVal);
3174            }
3175        }
3176        // Remove val from the result, by applyin inverse
3177        result = INTF::doInverse(result, val);
3178    }
3179    else
3180    {
3181        result = INTF::getInitial(val);
3182        if (!_waveIsSingleLane(mask))
3183        {
3184            int remaining = mask;
3185            while (remaining)
3186            {
3187                const int laneBit = remaining & -remaining;
3188                // Get the sourceLane
3189                const int srcLane = __ffs(laneBit) - 1;
3190                // Broadcast (can also broadcast to self)
3191                const T readValue = __shfl_sync(mask, val, srcLane);
3192                // Only accumulate if srcLane is less than this lane
3193                if (srcLane < laneId)
3194                {
3195                    result = INTF::doOp(result, readValue);
3196                }
3197                remaining &= ~laneBit;
3198            }
3199        }
3200    }
3201    return result;
3202}
3203
3204
3205// This implementation separately tracks the value to be propogated, and the value
3206// that is the final result
3207template<typename INTF, typename T>
3208__device__ T _wavePrefixScalar(WarpMask mask, T val)
3209{
3210    const int offsetSize = _waveCalcPow2Offset(mask);
3211
3212    const int laneId = _getLaneId();
3213    T result = INTF::getInitial(val);
3214    if (offsetSize > 0)
3215    {
3216        // For transmitted value we will do it inclusively with this lanes value
3217        // For the result we do not include the lanes value. This means an extra multiply for each
3218        // iteration but means we don't need to have a divide at the end and also removes overflow
3219        // issues in that scenario.
3220        for (int i = 1; i < offsetSize; i += i)
3221        {
3222            const T readVal = __shfl_up_sync(mask, val, i, offsetSize);
3223            if (laneId >= i)
3224            {
3225                result = INTF::doOp(result, readVal);
3226                val = INTF::doOp(val, readVal);
3227            }
3228        }
3229    }
3230    else
3231    {
3232        if (!_waveIsSingleLane(mask))
3233        {
3234            int remaining = mask;
3235            while (remaining)
3236            {
3237                const int laneBit = remaining & -remaining;
3238                // Get the sourceLane
3239                const int srcLane = __ffs(laneBit) - 1;
3240                // Broadcast (can also broadcast to self)
3241                const T readValue = __shfl_sync(mask, val, srcLane);
3242                // Only accumulate if srcLane is less than this lane
3243                if (srcLane < laneId)
3244                {
3245                    result = INTF::doOp(result, readValue);
3246                }
3247                remaining &= ~laneBit;
3248            }
3249        }
3250    }
3251    return result;
3252}
3253
3254
3255template<typename INTF, typename T, size_t COUNT>
3256__device__ T _waveOpCopy(T* dst, const T* src)
3257{
3258    for (size_t j = 0; j < COUNT; ++j)
3259    {
3260        dst[j] = src[j];
3261    }
3262}
3263
3264
3265template<typename INTF, typename T, size_t COUNT>
3266__device__ T _waveOpDoInverse(T* inOut, const T* val)
3267{
3268    for (size_t j = 0; j < COUNT; ++j)
3269    {
3270        inOut[j] = INTF::doInverse(inOut[j], val[j]);
3271    }
3272}
3273
3274template<typename INTF, typename T, size_t COUNT>
3275__device__ T _waveOpSetInitial(T* out, const T* val)
3276{
3277    for (size_t j = 0; j < COUNT; ++j)
3278    {
3279        out[j] = INTF::getInitial(val[j]);
3280    }
3281}
3282
3283template<typename INTF, typename T, size_t COUNT>
3284__device__ T _wavePrefixInvertableMultiple(WarpMask mask, T* val)
3285{
3286    const int offsetSize = _waveCalcPow2Offset(mask);
3287
3288    const int laneId = _getLaneId();
3289    T originalVal[COUNT];
3290    _waveOpCopy<INTF, T, COUNT>(originalVal, val);
3291
3292    if (offsetSize > 0)
3293    {
3294        // Sum is calculated inclusive of this lanes value
3295        for (int i = 1; i < offsetSize; i += i)
3296        {
3297            // TODO(JS): Note that here I don't split the laneId outside so it's only tested once.
3298            // This may be better but it would also mean that there would be shfl between lanes
3299            // that are on different (albeit identical) instructions. So this seems more likely to
3300            // work as expected with everything in lock step.
3301            for (size_t j = 0; j < COUNT; ++j)
3302            {
3303                const T readVal = __shfl_up_sync(mask, val[j], i, offsetSize);
3304                if (laneId >= i)
3305                {
3306                    val[j] = INTF::doOp(val[j], readVal);
3307                }
3308            }
3309        }
3310        // Remove originalVal from the result, by applyin inverse
3311        _waveOpDoInverse<INTF, T, COUNT>(val, originalVal);
3312    }
3313    else
3314    {
3315        _waveOpSetInitial<INTF, T, COUNT>(val, val);
3316        if (!_waveIsSingleLane(mask))
3317        {
3318            int remaining = mask;
3319            while (remaining)
3320            {
3321                const int laneBit = remaining & -remaining;
3322                // Get the sourceLane
3323                const int srcLane = __ffs(laneBit) - 1;
3324
3325                for (size_t j = 0; j < COUNT; ++j)
3326                {
3327                    // Broadcast (can also broadcast to self)
3328                    const T readValue = __shfl_sync(mask, originalVal[j], srcLane);
3329                    // Only accumulate if srcLane is less than this lane
3330                    if (srcLane < laneId)
3331                    {
3332                        val[j] = INTF::doOp(val[j], readValue);
3333                    }
3334                    remaining &= ~laneBit;
3335                }
3336            }
3337        }
3338    }
3339}
3340
3341template<typename INTF, typename T, size_t COUNT>
3342__device__ T _wavePrefixMultiple(WarpMask mask, T* val)
3343{
3344    const int offsetSize = _waveCalcPow2Offset(mask);
3345
3346    const int laneId = _getLaneId();
3347
3348    T work[COUNT];
3349    _waveOpCopy<INTF, T, COUNT>(work, val);
3350    _waveOpSetInitial<INTF, T, COUNT>(val, val);
3351
3352    if (offsetSize > 0)
3353    {
3354        // For transmitted value we will do it inclusively with this lanes value
3355        // For the result we do not include the lanes value. This means an extra op for each
3356        // iteration but means we don't need to have a divide at the end and also removes overflow
3357        // issues in that scenario.
3358        for (int i = 1; i < offsetSize; i += i)
3359        {
3360            for (size_t j = 0; j < COUNT; ++j)
3361            {
3362                const T readVal = __shfl_up_sync(mask, work[j], i, offsetSize);
3363                if (laneId >= i)
3364                {
3365                    work[j] = INTF::doOp(work[j], readVal);
3366                    val[j] = INTF::doOp(val[j], readVal);
3367                }
3368            }
3369        }
3370    }
3371    else
3372    {
3373        if (!_waveIsSingleLane(mask))
3374        {
3375            int remaining = mask;
3376            while (remaining)
3377            {
3378                const int laneBit = remaining & -remaining;
3379                // Get the sourceLane
3380                const int srcLane = __ffs(laneBit) - 1;
3381
3382                for (size_t j = 0; j < COUNT; ++j)
3383                {
3384                    // Broadcast (can also broadcast to self)
3385                    const T readValue = __shfl_sync(mask, work[j], srcLane);
3386                    // Only accumulate if srcLane is less than this lane
3387                    if (srcLane < laneId)
3388                    {
3389                        val[j] = INTF::doOp(val[j], readValue);
3390                    }
3391                }
3392                remaining &= ~laneBit;
3393            }
3394        }
3395    }
3396}
3397
3398template<typename T>
3399__inline__ __device__ T _wavePrefixProduct(WarpMask mask, T val)
3400{
3401    return _wavePrefixScalar<WaveOpMul<T>, T>(mask, val);
3402}
3403
3404template<typename T>
3405__inline__ __device__ T _wavePrefixSum(WarpMask mask, T val)
3406{
3407    return _wavePrefixInvertableScalar<WaveOpAdd<T>, T>(mask, val);
3408}
3409
3410template<typename T>
3411__inline__ __device__ T _wavePrefixXor(WarpMask mask, T val)
3412{
3413    return _wavePrefixInvertableScalar<WaveOpXor<T>, T>(mask, val);
3414}
3415
3416template<typename T>
3417__inline__ __device__ T _wavePrefixOr(WarpMask mask, T val)
3418{
3419    return _wavePrefixScalar<WaveOpOr<T>, T>(mask, val);
3420}
3421
3422template<typename T>
3423__inline__ __device__ T _wavePrefixAnd(WarpMask mask, T val)
3424{
3425    return _wavePrefixScalar<WaveOpAnd<T>, T>(mask, val);
3426}
3427
3428
3429template<typename T>
3430__inline__ __device__ T _wavePrefixProductMultiple(WarpMask mask, T val)
3431{
3432    typedef typename ElementTypeTrait<T>::Type ElemType;
3433    _wavePrefixInvertableMultiple<WaveOpMul<ElemType>, ElemType, sizeof(T) / sizeof(ElemType)>(
3434        mask,
3435        (ElemType*)&val);
3436    return val;
3437}
3438
3439template<typename T>
3440__inline__ __device__ T _wavePrefixSumMultiple(WarpMask mask, T val)
3441{
3442    typedef typename ElementTypeTrait<T>::Type ElemType;
3443    _wavePrefixInvertableMultiple<WaveOpAdd<ElemType>, ElemType, sizeof(T) / sizeof(ElemType)>(
3444        mask,
3445        (ElemType*)&val);
3446    return val;
3447}
3448
3449template<typename T>
3450__inline__ __device__ T _wavePrefixXorMultiple(WarpMask mask, T val)
3451{
3452    typedef typename ElementTypeTrait<T>::Type ElemType;
3453    _wavePrefixInvertableMultiple<WaveOpXor<ElemType>, ElemType, sizeof(T) / sizeof(ElemType)>(
3454        mask,
3455        (ElemType*)&val);
3456    return val;
3457}
3458
3459template<typename T>
3460__inline__ __device__ T _wavePrefixOrMultiple(WarpMask mask, T val)
3461{
3462    typedef typename ElementTypeTrait<T>::Type ElemType;
3463    _wavePrefixMultiple<WaveOpOr<ElemType>, ElemType, sizeof(T) / sizeof(ElemType)>(
3464        mask,
3465        (ElemType*)&val);
3466    return val;
3467}
3468
3469template<typename T>
3470__inline__ __device__ T _wavePrefixAndMultiple(WarpMask mask, T val)
3471{
3472    typedef typename ElementTypeTrait<T>::Type ElemType;
3473    _wavePrefixMultiple<WaveOpAnd<ElemType>, ElemType, sizeof(T) / sizeof(ElemType)>(
3474        mask,
3475        (ElemType*)&val);
3476    return val;
3477}
3478
3479template<typename T>
3480__inline__ __device__ T _wavePrefixMin(WarpMask mask, T val)
3481{
3482    return _wavePrefixScalar<WaveOpMin<T>, T>(mask, val);
3483}
3484
3485template<typename T>
3486__inline__ __device__ T _wavePrefixMax(WarpMask mask, T val)
3487{
3488    return _wavePrefixScalar<WaveOpMax<T>, T>(mask, val);
3489}
3490
3491template<typename T>
3492__inline__ __device__ T _wavePrefixMinMultiple(WarpMask mask, T val)
3493{
3494    typedef typename ElementTypeTrait<T>::Type ElemType;
3495    _wavePrefixMultiple<WaveOpMin<ElemType>, ElemType, sizeof(T) / sizeof(ElemType)>(
3496        mask,
3497        (ElemType*)&val);
3498    return val;
3499}
3500
3501template<typename T>
3502__inline__ __device__ T _wavePrefixMaxMultiple(WarpMask mask, T val)
3503{
3504    typedef typename ElementTypeTrait<T>::Type ElemType;
3505    _wavePrefixMultiple<WaveOpMax<ElemType>, ElemType, sizeof(T) / sizeof(ElemType)>(
3506        mask,
3507        (ElemType*)&val);
3508    return val;
3509}
3510
3511// Wrapper structures for exclusive operations that use the overloaded getInitial method
3512template<typename T>
3513struct WaveOpExclusiveMin
3514{
3515    __inline__ __device__ static T getInitial(T a) { return WaveOpMin<T>::getInitial(a, true); }
3516    __inline__ __device__ static T doOp(T a, T b) { return WaveOpMin<T>::doOp(a, b); }
3517};
3518
3519template<typename T>
3520struct WaveOpExclusiveMax
3521{
3522    __inline__ __device__ static T getInitial(T a) { return WaveOpMax<T>::getInitial(a, true); }
3523    __inline__ __device__ static T doOp(T a, T b) { return WaveOpMax<T>::doOp(a, b); }
3524};
3525
3526// Inclusive prefix min/max functions (for WaveMultiPrefixInclusive*)
3527template<typename T>
3528__inline__ __device__ T _wavePrefixInclusiveMin(WarpMask mask, T val)
3529{
3530    return _wavePrefixMin(mask, val);
3531}
3532
3533template<typename T>
3534__inline__ __device__ T _wavePrefixInclusiveMax(WarpMask mask, T val)
3535{
3536    return _wavePrefixMax(mask, val);
3537}
3538
3539template<typename T>
3540__inline__ __device__ T _wavePrefixInclusiveMinMultiple(WarpMask mask, T val)
3541{
3542    return _wavePrefixMinMultiple(mask, val);
3543}
3544
3545template<typename T>
3546__inline__ __device__ T _wavePrefixInclusiveMaxMultiple(WarpMask mask, T val)
3547{
3548    return _wavePrefixMaxMultiple(mask, val);
3549}
3550
3551// Explicit exclusive prefix min/max functions (for WaveMultiPrefixExclusive*)
3552template<typename T>
3553__inline__ __device__ T _wavePrefixExclusiveMin(WarpMask mask, T val)
3554{
3555    return _wavePrefixScalar<WaveOpExclusiveMin<T>, T>(mask, val);
3556}
3557
3558template<typename T>
3559__inline__ __device__ T _wavePrefixExclusiveMax(WarpMask mask, T val)
3560{
3561    return _wavePrefixScalar<WaveOpExclusiveMax<T>, T>(mask, val);
3562}
3563
3564template<typename T>
3565__inline__ __device__ T _wavePrefixExclusiveMinMultiple(WarpMask mask, T val)
3566{
3567    typedef typename ElementTypeTrait<T>::Type ElemType;
3568    _wavePrefixMultiple<WaveOpExclusiveMin<ElemType>, ElemType, sizeof(T) / sizeof(ElemType)>(
3569        mask,
3570        (ElemType*)&val);
3571    return val;
3572}
3573
3574template<typename T>
3575__inline__ __device__ T _wavePrefixExclusiveMaxMultiple(WarpMask mask, T val)
3576{
3577    typedef typename ElementTypeTrait<T>::Type ElemType;
3578    _wavePrefixMultiple<WaveOpExclusiveMax<ElemType>, ElemType, sizeof(T) / sizeof(ElemType)>(
3579        mask,
3580        (ElemType*)&val);
3581    return val;
3582}
3583
3584template<typename T>
3585__inline__ __device__ uint4 _waveMatchScalar(WarpMask mask, T val)
3586{
3587    int pred;
3588    return make_uint4(__match_all_sync(mask, val, &pred), 0, 0, 0);
3589}
3590
3591template<typename T>
3592__inline__ __device__ uint4 _waveMatchMultiple(WarpMask mask, const T& inVal)
3593{
3594    typedef typename ElementTypeTrait<T>::Type ElemType;
3595    const size_t count = sizeof(T) / sizeof(ElemType);
3596    int pred;
3597    const ElemType* src = (const ElemType*)&inVal;
3598    uint matchBits = 0xffffffff;
3599    for (size_t i = 0; i < count && matchBits; ++i)
3600    {
3601        matchBits = matchBits & __match_all_sync(mask, src[i], &pred);
3602    }
3603    return make_uint4(matchBits, 0, 0, 0);
3604}
3605
3606__device__ uint getAt(dim3 a, int b)
3607{
3608    SLANG_PRELUDE_ASSERT(b >= 0 && b < 3);
3609    return (&a.x)[b];
3610}
3611__device__ uint3 operator*(uint3 a, dim3 b)
3612{
3613    uint3 r;
3614    r.x = a.x * b.x;
3615    r.y = a.y * b.y;
3616    r.z = a.z * b.z;
3617    return r;
3618}
3619
3620template<typename TResult, typename TInput>
3621__inline__ __device__ TResult slang_bit_cast(TInput val)
3622{
3623    return *(TResult*)(&val);
3624}
3625
3626/* !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! */
3627
3628
3629/* Type that defines the uniform entry point params. The actual content of this type is dependent on
3630the entry point parameters, and can be found via reflection or defined such that it matches the
3631shader appropriately.
3632*/
3633struct UniformEntryPointParams;
3634struct UniformState;
3635
3636// ---------------------- OptiX Ray Payload --------------------------------------
3637#ifdef SLANG_CUDA_ENABLE_OPTIX
3638
3639struct RayDesc
3640{
3641    float3 Origin;
3642    float TMin;
3643    float3 Direction;
3644    float TMax;
3645};
3646
3647static __forceinline__ __device__ void* unpackOptiXRayPayloadPointer(uint32_t i0, uint32_t i1)
3648{
3649    const uint64_t uptr = static_cast<uint64_t>(i0) << 32 | i1;
3650    void* ptr = reinterpret_cast<void*>(uptr);
3651    return ptr;
3652}
3653
3654static __forceinline__ __device__ void packOptiXRayPayloadPointer(
3655    void* ptr,
3656    uint32_t& i0,
3657    uint32_t& i1)
3658{
3659    const uint64_t uptr = reinterpret_cast<uint64_t>(ptr);
3660    i0 = uptr >> 32;
3661    i1 = uptr & 0x00000000ffffffff;
3662}
3663
3664static __forceinline__ __device__ void* getOptiXRayPayloadPtr()
3665{
3666    const uint32_t u0 = optixGetPayload_0();
3667    const uint32_t u1 = optixGetPayload_1();
3668    return unpackOptiXRayPayloadPointer(u0, u1);
3669}
3670
3671template<typename T>
3672__forceinline__ __device__ void* optixTrace(
3673    OptixTraversableHandle AccelerationStructure,
3674    uint32_t RayFlags,
3675    uint32_t InstanceInclusionMask,
3676    uint32_t RayContributionToHitGroupIndex,
3677    uint32_t MultiplierForGeometryContributionToHitGroupIndex,
3678    uint32_t MissShaderIndex,
3679    RayDesc Ray,
3680    T* Payload)
3681{
3682    uint32_t r0, r1;
3683    packOptiXRayPayloadPointer((void*)Payload, r0, r1);
3684    optixTrace(
3685        AccelerationStructure,
3686        Ray.Origin,
3687        Ray.Direction,
3688        Ray.TMin,
3689        Ray.TMax,
3690        0.f, /* Time for motion blur, currently unsupported in slang */
3691        InstanceInclusionMask,
3692        RayFlags,
3693        RayContributionToHitGroupIndex,
3694        MultiplierForGeometryContributionToHitGroupIndex,
3695        MissShaderIndex,
3696        r0,
3697        r1);
3698}
3699
3700#if (OPTIX_VERSION >= 90000)
3701__forceinline__ __device__ float4 optixGetSpherePositionAndRadius()
3702{
3703    float4 data[1];
3704    optixGetSphereData(data);
3705    return data[0];
3706}
3707#endif
3708
3709#if (OPTIX_VERSION >= 90000)
3710__forceinline__ __device__ float4
3711optixHitObjectGetSpherePositionAndRadius(OptixTraversableHandle* Obj)
3712{
3713    float4 data[1];
3714    optixHitObjectGetSphereData(data);
3715    return data[0];
3716}
3717#endif
3718
3719#if (OPTIX_VERSION >= 90000)
3720__forceinline__ __device__ Matrix<float, 2, 4> optixGetLssPositionsAndRadii()
3721{
3722    float4 data[2];
3723    optixGetLinearCurveVertexData(data);
3724    return makeMatrix<float, 2, 4>(data[0], data[1]);
3725}
3726#endif
3727
3728#if (OPTIX_VERSION >= 90000)
3729__forceinline__ __device__ Matrix<float, 2, 4> optixHitObjectGetLssPositionsAndRadii(
3730    OptixTraversableHandle* Obj)
3731{
3732    float4 data[2];
3733    optixHitObjectGetLinearCurveVertexData(data);
3734    return makeMatrix<float, 2, 4>(data[0], data[1]);
3735}
3736#endif
3737
3738#if (OPTIX_VERSION >= 90000)
3739__forceinline__ __device__ bool optixIsSphereHit()
3740{
3741    return optixGetPrimitiveType() == OPTIX_PRIMITIVE_TYPE_SPHERE;
3742}
3743#endif
3744
3745#if (OPTIX_VERSION >= 90000)
3746__forceinline__ __device__ bool optixHitObjectIsSphereHit(OptixTraversableHandle* Obj)
3747{
3748    return optixGetPrimitiveType(optixHitObjectGetHitKind()) == OPTIX_PRIMITIVE_TYPE_SPHERE;
3749}
3750#endif
3751
3752#if (OPTIX_VERSION >= 90000)
3753__forceinline__ __device__ bool optixIsLSSHit()
3754{
3755    return optixGetPrimitiveType() == OPTIX_PRIMITIVE_TYPE_ROUND_LINEAR;
3756}
3757#endif
3758
3759#if (OPTIX_VERSION >= 90000)
3760__forceinline__ __device__ bool optixHitObjectIsLSSHit(OptixTraversableHandle* Obj)
3761{
3762    return optixGetPrimitiveType(optixHitObjectGetHitKind()) == OPTIX_PRIMITIVE_TYPE_ROUND_LINEAR;
3763}
3764#endif
3765
3766template<typename T>
3767__forceinline__ __device__ void* optixTraverse(
3768    OptixTraversableHandle AccelerationStructure,
3769    uint32_t RayFlags,
3770    uint32_t InstanceInclusionMask,
3771    uint32_t RayContributionToHitGroupIndex,
3772    uint32_t MultiplierForGeometryContributionToHitGroupIndex,
3773    uint32_t MissShaderIndex,
3774    RayDesc Ray,
3775    T* Payload,
3776    OptixTraversableHandle* hitObj)
3777{
3778    uint32_t r0, r1;
3779    packOptiXRayPayloadPointer((void*)Payload, r0, r1);
3780    optixTraverse(
3781        AccelerationStructure,
3782        Ray.Origin,
3783        Ray.Direction,
3784        Ray.TMin,
3785        Ray.TMax,
3786        0.f, /* Time for motion blur, currently unsupported in slang */
3787        InstanceInclusionMask,
3788        RayFlags,
3789        RayContributionToHitGroupIndex,
3790        MultiplierForGeometryContributionToHitGroupIndex,
3791        MissShaderIndex,
3792        r0,
3793        r1);
3794}
3795
3796template<typename T>
3797__forceinline__ __device__ void* optixTraverse(
3798    OptixTraversableHandle AccelerationStructure,
3799    uint32_t RayFlags,
3800    uint32_t InstanceInclusionMask,
3801    uint32_t RayContributionToHitGroupIndex,
3802    uint32_t MultiplierForGeometryContributionToHitGroupIndex,
3803    uint32_t MissShaderIndex,
3804    RayDesc Ray,
3805    float RayTime,
3806    T* Payload,
3807    OptixTraversableHandle* hitObj)
3808{
3809    uint32_t r0, r1;
3810    packOptiXRayPayloadPointer((void*)Payload, r0, r1);
3811    optixTraverse(
3812        AccelerationStructure,
3813        Ray.Origin,
3814        Ray.Direction,
3815        Ray.TMin,
3816        Ray.TMax,
3817        RayTime,
3818        InstanceInclusionMask,
3819        RayFlags,
3820        RayContributionToHitGroupIndex,
3821        MultiplierForGeometryContributionToHitGroupIndex,
3822        MissShaderIndex,
3823        r0,
3824        r1);
3825}
3826
3827#if (OPTIX_VERSION >= 80100)
3828static __forceinline__ __device__ bool slangOptixHitObjectIsHit(OptixTraversableHandle* hitObj)
3829{
3830    return optixHitObjectIsHit();
3831}
3832#endif
3833
3834#if (OPTIX_VERSION >= 80100)
3835static __forceinline__ __device__ bool slangOptixHitObjectIsMiss(OptixTraversableHandle* hitObj)
3836{
3837    return optixHitObjectIsMiss();
3838}
3839#endif
3840
3841#if (OPTIX_VERSION >= 80100)
3842static __forceinline__ __device__ bool slangOptixHitObjectIsNop(OptixTraversableHandle* hitObj)
3843{
3844    return optixHitObjectIsNop();
3845}
3846#endif
3847
3848#if (OPTIX_VERSION >= 90000)
3849static __forceinline__ __device__ uint
3850slangOptixHitObjectGetClusterId(OptixTraversableHandle* hitObj)
3851{
3852    return optixHitObjectGetClusterId();
3853}
3854#endif
3855
3856#if (OPTIX_VERSION >= 80100)
3857static __forceinline__ __device__ void optixMakeMissHitObject(
3858    uint MissShaderIndex,
3859    RayDesc Ray,
3860    OptixTraversableHandle* missObj)
3861{
3862    optixMakeMissHitObject(
3863        MissShaderIndex,
3864        Ray.Origin,
3865        Ray.Direction,
3866        Ray.TMin,
3867        Ray.TMax,
3868        0.f /* rayTime */
3869#if (OPTIX_VERSION >= 90000)
3870        ,
3871        OPTIX_RAY_FLAG_NONE /* rayFlags*/
3872#endif
3873    );
3874}
3875#endif
3876
3877#if (OPTIX_VERSION >= 80100)
3878static __forceinline__ __device__ void optixMakeMissHitObject(
3879    uint MissShaderIndex,
3880    RayDesc Ray,
3881    float CurrentTime,
3882    OptixTraversableHandle* missObj)
3883{
3884    optixMakeMissHitObject(
3885        MissShaderIndex,
3886        Ray.Origin,
3887        Ray.Direction,
3888        Ray.TMin,
3889        Ray.TMax,
3890        CurrentTime
3891#if (OPTIX_VERSION >= 90000)
3892        ,
3893        OPTIX_RAY_FLAG_NONE /* rayFlags*/
3894#endif
3895    );
3896}
3897#endif
3898
3899#if (OPTIX_VERSION >= 80100)
3900template<typename T>
3901static __forceinline__ __device__ void optixMakeHitObject(
3902    OptixTraversableHandle AccelerationStructure,
3903    uint InstanceIndex,
3904    uint GeometryIndex,
3905    uint PrimitiveIndex,
3906    uint HitKind,
3907    uint RayContributionToHitGroupIndex,
3908    uint MultiplierForGeometryContributionToHitGroupIndex,
3909    RayDesc Ray,
3910    T attr,
3911    OptixTraversableHandle* handle)
3912{
3913    OptixTraverseData data{};
3914    optixHitObjectGetTraverseData(&data);
3915    optixMakeHitObject(
3916        AccelerationStructure,
3917        Ray.Origin,
3918        Ray.Direction,
3919        Ray.TMin,
3920        0.f,
3921        OPTIX_RAY_FLAG_NONE, /* rayFlags*/
3922        data,
3923        nullptr, /*OptixTraversableHandle* transforms*/
3924        0 /*numTransforms */);
3925}
3926#endif
3927
3928#if (OPTIX_VERSION >= 80100)
3929template<typename T>
3930static __forceinline__ __device__ void optixMakeHitObject(
3931    uint HitGroupRecordIndex,
3932    OptixTraversableHandle AccelerationStructure,
3933    uint InstanceIndex,
3934    uint GeometryIndex,
3935    uint PrimitiveIndex,
3936    uint HitKind,
3937    RayDesc Ray,
3938    T attr,
3939    OptixTraversableHandle* handle)
3940{
3941    OptixTraverseData data{};
3942    optixHitObjectGetTraverseData(&data);
3943    optixMakeHitObject(
3944        AccelerationStructure,
3945        Ray.Origin,
3946        Ray.Direction,
3947        Ray.TMin,
3948        0.f,
3949        OPTIX_RAY_FLAG_NONE, /* rayFlags*/
3950        data,
3951        nullptr, /*OptixTraversableHandle* transforms*/
3952        0 /*numTransforms */);
3953}
3954#endif
3955
3956#if (OPTIX_VERSION >= 80100)
3957template<typename T>
3958static __forceinline__ __device__ void optixMakeHitObject(
3959    OptixTraversableHandle AccelerationStructure,
3960    uint InstanceIndex,
3961    uint GeometryIndex,
3962    uint PrimitiveIndex,
3963    uint HitKind,
3964    uint RayContributionToHitGroupIndex,
3965    uint MultiplierForGeometryContributionToHitGroupIndex,
3966    RayDesc Ray,
3967    float CurrentTime,
3968    T attr,
3969    OptixTraversableHandle* handle)
3970{
3971    OptixTraverseData data{};
3972    optixHitObjectGetTraverseData(&data);
3973    optixMakeHitObject(
3974        AccelerationStructure,
3975        Ray.Origin,
3976        Ray.Direction,
3977        Ray.TMin,
3978        CurrentTime,
3979        OPTIX_RAY_FLAG_NONE, /* rayFlags*/
3980        data,
3981        nullptr, /*OptixTraversableHandle* transforms*/
3982        0 /*numTransforms */);
3983}
3984#endif
3985
3986#if (OPTIX_VERSION >= 80100)
3987template<typename T>
3988static __forceinline__ __device__ void optixMakeHitObject(
3989    uint HitGroupRecordIndex,
3990    OptixTraversableHandle AccelerationStructure,
3991    uint InstanceIndex,
3992    uint GeometryIndex,
3993    uint PrimitiveIndex,
3994    uint HitKind,
3995    RayDesc Ray,
3996    float CurrentTime,
3997    T attr,
3998    OptixTraversableHandle* handle)
3999{
4000    OptixTraverseData data{};
4001    optixHitObjectGetTraverseData(&data);
4002    optixMakeHitObject(
4003        AccelerationStructure,
4004        Ray.Origin,
4005        Ray.Direction,
4006        Ray.TMin,
4007        CurrentTime,
4008        OPTIX_RAY_FLAG_NONE, /* rayFlags*/
4009        data,
4010        nullptr, /*OptixTraversableHandle* transforms*/
4011        0 /*numTransforms */);
4012}
4013#endif
4014
4015#if (OPTIX_VERSION >= 80100)
4016static __forceinline__ __device__ void slangOptixMakeNopHitObject(OptixTraversableHandle* Obj)
4017{
4018    optixMakeNopHitObject();
4019}
4020#endif
4021
4022#if (OPTIX_VERSION >= 80100)
4023template<typename T>
4024static __forceinline__ __device__ void optixInvoke(
4025    OptixTraversableHandle AccelerationStructure,
4026    OptixTraversableHandle* HitOrMiss,
4027    T Payload)
4028{
4029    uint32_t r0, r1;
4030    packOptiXRayPayloadPointer((void*)Payload, r0, r1);
4031    optixInvoke(r0, r1);
4032}
4033#endif
4034
4035#if (OPTIX_VERSION >= 80100)
4036static __forceinline__ __device__ RayDesc optixHitObjectGetRayDesc(OptixTraversableHandle* obj)
4037{
4038    RayDesc ray = {
4039        optixHitObjectGetWorldRayOrigin(),
4040        optixHitObjectGetRayTmin(),
4041        optixHitObjectGetWorldRayDirection(),
4042        optixHitObjectGetRayTmax()};
4043    return ray;
4044}
4045#endif
4046
4047#if (OPTIX_VERSION >= 80100)
4048static __forceinline__ __device__ uint
4049slangOptixHitObjectGetInstanceIndex(OptixTraversableHandle* Obj)
4050{
4051    return optixHitObjectGetInstanceIndex();
4052}
4053#endif
4054
4055#if (OPTIX_VERSION >= 80100)
4056static __forceinline__ __device__ uint slangOptixHitObjectGetInstanceId(OptixTraversableHandle* Obj)
4057{
4058    return optixHitObjectGetInstanceId();
4059}
4060#endif
4061
4062#if (OPTIX_VERSION >= 80100)
4063static __forceinline__ __device__ uint
4064slangOptixHitObjectGetSbtGASIndex(OptixTraversableHandle* Obj)
4065{
4066    return optixHitObjectGetSbtGASIndex();
4067}
4068#endif
4069
4070#if (OPTIX_VERSION >= 80100)
4071static __forceinline__ __device__ uint
4072slangOptixHitObjectGetPrimitiveIndex(OptixTraversableHandle* Obj)
4073{
4074    return optixHitObjectGetPrimitiveIndex();
4075}
4076#endif
4077
4078#if (OPTIX_VERSION >= 80100)
4079template<typename T>
4080static __forceinline__ __device__ T optixHitObjectGetAttribute(OptixTraversableHandle* Obj)
4081{
4082    constexpr size_t numInts = (sizeof(T) + sizeof(uint32_t) - 1) /
4083                               sizeof(uint32_t); // Number of 32-bit values, rounded up
4084    static_assert(numInts <= 8, "Attribute type is too large");
4085
4086    // Create an array to hold the attribute values
4087    uint32_t values[numInts == 0 ? 1 : numInts] = {0}; // Ensure we have at least one element
4088
4089    // Read the appropriate number of attribute registers
4090    if constexpr (numInts > 0)
4091        values[0] = optixHitObjectGetAttribute_0();
4092    if constexpr (numInts > 1)
4093        values[1] = optixHitObjectGetAttribute_1();
4094    if constexpr (numInts > 2)
4095        values[2] = optixHitObjectGetAttribute_2();
4096    if constexpr (numInts > 3)
4097        values[3] = optixHitObjectGetAttribute_3();
4098    if constexpr (numInts > 4)
4099        values[4] = optixHitObjectGetAttribute_4();
4100    if constexpr (numInts > 5)
4101        values[5] = optixHitObjectGetAttribute_5();
4102    if constexpr (numInts > 6)
4103        values[6] = optixHitObjectGetAttribute_6();
4104    if constexpr (numInts > 7)
4105        values[7] = optixHitObjectGetAttribute_7();
4106
4107    // Reinterpret the array as the desired type
4108    T result;
4109    memcpy(&result, values, sizeof(T));
4110    return result;
4111}
4112#endif
4113
4114#if (OPTIX_VERSION >= 80100)
4115static __forceinline__ __device__ uint
4116slangOptixHitObjectGetSbtRecordIndex(OptixTraversableHandle* Obj)
4117{
4118    return optixHitObjectGetSbtRecordIndex();
4119}
4120#endif
4121
4122#if (OPTIX_VERSION >= 90000)
4123static __forceinline__ __device__ uint
4124slangOptixHitObjectSetSbtRecordIndex(OptixTraversableHandle* Obj, uint sbtRecordIndex)
4125{
4126    optixHitObjectSetSbtRecordIndex(sbtRecordIndex); // returns void
4127    return sbtRecordIndex;
4128}
4129#endif
4130
4131#else
4132// Define OptixTraversableHandle even if OptiX is not enabled.
4133// This allows RaytracingAccelerationStructure to be properly reflected in non-OptiX code.
4134typedef unsigned long long OptixTraversableHandle;
4135#endif
4136static const int kSlangTorchTensorMaxDim = 5;
4137
4138// TensorView
4139struct TensorView
4140{
4141    uint8_t* data;
4142    uint32_t strides[kSlangTorchTensorMaxDim];
4143    uint32_t sizes[kSlangTorchTensorMaxDim];
4144    uint32_t dimensionCount;
4145
4146    template<typename T>
4147    __device__ T* data_ptr()
4148    {
4149        return reinterpret_cast<T*>(data);
4150    }
4151
4152    template<typename T>
4153    __device__ T* data_ptr_at(uint32_t index)
4154    {
4155        uint64_t offset = strides[0] * index;
4156        return reinterpret_cast<T*>(data + offset);
4157    }
4158
4159    template<typename T>
4160    __device__ T* data_ptr_at(uint2 index)
4161    {
4162        uint64_t offset = strides[0] * index.x + strides[1] * index.y;
4163        return reinterpret_cast<T*>(data + offset);
4164    }
4165
4166    template<typename T>
4167    __device__ T* data_ptr_at(uint3 index)
4168    {
4169        uint64_t offset = strides[0] * index.x + strides[1] * index.y + strides[2] * index.z;
4170        return reinterpret_cast<T*>(data + offset);
4171    }
4172
4173    template<typename T>
4174    __device__ T* data_ptr_at(uint4 index)
4175    {
4176        uint64_t offset = strides[0] * index.x + strides[1] * index.y + strides[2] * index.z +
4177                          strides[3] * index.w;
4178        return reinterpret_cast<T*>(data + offset);
4179    }
4180
4181    template<typename T, unsigned int N>
4182    __device__ T* data_ptr_at(uint index[N])
4183    {
4184        uint64_t offset = 0;
4185        for (unsigned int i = 0; i < N; ++i)
4186        {
4187            offset += strides[i] * index[i];
4188        }
4189        return reinterpret_cast<T*>(data + offset);
4190    }
4191
4192    template<typename T>
4193    __device__ T& load(uint32_t x)
4194    {
4195        return *reinterpret_cast<T*>(data + strides[0] * x);
4196    }
4197    template<typename T>
4198    __device__ T& load(uint32_t x, uint32_t y)
4199    {
4200        return *reinterpret_cast<T*>(data + strides[0] * x + strides[1] * y);
4201    }
4202    template<typename T>
4203    __device__ T& load(uint2 index)
4204    {
4205        return *reinterpret_cast<T*>(data + strides[0] * index.x + strides[1] * index.y);
4206    }
4207    template<typename T>
4208    __device__ T& load(uint32_t x, uint32_t y, uint32_t z)
4209    {
4210        return *reinterpret_cast<T*>(data + strides[0] * x + strides[1] * y + strides[2] * z);
4211    }
4212    template<typename T>
4213    __device__ T& load(uint3 index)
4214    {
4215        return *reinterpret_cast<T*>(
4216            data + strides[0] * index.x + strides[1] * index.y + strides[2] * index.z);
4217    }
4218    template<typename T>
4219    __device__ T& load(uint32_t x, uint32_t y, uint32_t z, uint32_t w)
4220    {
4221        return *reinterpret_cast<T*>(
4222            data + strides[0] * x + strides[1] * y + strides[2] * z + strides[3] * w);
4223    }
4224    template<typename T>
4225    __device__ T& load(uint4 index)
4226    {
4227        return *reinterpret_cast<T*>(
4228            data + strides[0] * index.x + strides[1] * index.y + strides[2] * index.z +
4229            strides[3] * index.w);
4230    }
4231    template<typename T>
4232    __device__ T& load(uint32_t i0, uint32_t i1, uint32_t i2, uint32_t i3, uint32_t i4)
4233    {
4234        return *reinterpret_cast<T*>(
4235            data + strides[0] * i0 + strides[1] * i1 + strides[2] * i2 + strides[3] * i3 +
4236            strides[4] * i4);
4237    }
4238
4239    // Generic version of load
4240    template<typename T, unsigned int N>
4241    __device__ T& load(uint index[N])
4242    {
4243        uint64_t offset = 0;
4244        for (unsigned int i = 0; i < N; ++i)
4245        {
4246            offset += strides[i] * index[i];
4247        }
4248        return *reinterpret_cast<T*>(data + offset);
4249    }
4250
4251    template<typename T>
4252    __device__ void store(uint32_t x, T val)
4253    {
4254        *reinterpret_cast<T*>(data + strides[0] * x) = val;
4255    }
4256    template<typename T>
4257    __device__ void store(uint32_t x, uint32_t y, T val)
4258    {
4259        *reinterpret_cast<T*>(data + strides[0] * x + strides[1] * y) = val;
4260    }
4261    template<typename T>
4262    __device__ void store(uint2 index, T val)
4263    {
4264        *reinterpret_cast<T*>(data + strides[0] * index.x + strides[1] * index.y) = val;
4265    }
4266    template<typename T>
4267    __device__ void store(uint32_t x, uint32_t y, uint32_t z, T val)
4268    {
4269        *reinterpret_cast<T*>(data + strides[0] * x + strides[1] * y + strides[2] * z) = val;
4270    }
4271    template<typename T>
4272    __device__ void store(uint3 index, T val)
4273    {
4274        *reinterpret_cast<T*>(
4275            data + strides[0] * index.x + strides[1] * index.y + strides[2] * index.z) = val;
4276    }
4277    template<typename T>
4278    __device__ void store(uint32_t x, uint32_t y, uint32_t z, uint32_t w, T val)
4279    {
4280        *reinterpret_cast<T*>(
4281            data + strides[0] * x + strides[1] * y + strides[2] * z + strides[3] * w) = val;
4282    }
4283    template<typename T>
4284    __device__ void store(uint4 index, T val)
4285    {
4286        *reinterpret_cast<T*>(
4287            data + strides[0] * index.x + strides[1] * index.y + strides[2] * index.z +
4288            strides[3] * index.w) = val;
4289    }
4290    template<typename T>
4291    __device__ void store(uint32_t i0, uint32_t i1, uint32_t i2, uint32_t i3, uint32_t i4, T val)
4292    {
4293        *reinterpret_cast<T*>(
4294            data + strides[0] * i0 + strides[1] * i1 + strides[2] * i2 + strides[3] * i3 +
4295            strides[4] * i4) = val;
4296    }
4297
4298    // Generic version
4299    template<typename T, unsigned int N>
4300    __device__ void store(uint index[N], T val)
4301    {
4302        uint64_t offset = 0;
4303        for (unsigned int i = 0; i < N; ++i)
4304        {
4305            offset += strides[i] * index[i];
4306        }
4307        *reinterpret_cast<T*>(data + offset) = val;
4308    }
4309};
4310
4311// Implementations for texture fetch/load functions using tex PTX intrinsics
4312// These are used for read-only texture access with integer coordinates.
4313
4314// 1D is not supported via PTX. Keeping the implementation below in case it ever gets supported.
4315template<typename T>
4316SLANG_FORCE_INLINE SLANG_CUDA_CALL T tex1Dfetch_int(CUtexObject texObj, int x, int mip)
4317{
4318    static_assert(false, "CUDA does not support fetching from 1D textures");
4319}
4320
4321#if 0
4322#define SLANG_TEX1DFETCH_INT_IMPL(T, dtype, c)                                                 \
4323    template<>                                                                                 \
4324    SLANG_FORCE_INLINE SLANG_CUDA_CALL T tex1Dfetch_int(CUtexObject texObj, int x, int mip)    \
4325    {                                                                                          \
4326        T result;                                                                              \
4327        T stub;                                                                                \
4328        asm("tex.level.1d.v4." dtype ".s32 {%0, %1, %2, %3}, [%4, {%5}], %6;"                  \
4329            : c(result), c(stub), c(stub), c(stub)                                             \
4330            : "l"(texObj), "r"(x), "r"(mip));                                                  \
4331        return result;                                                                         \
4332    }                                                                                          \
4333    template<>                                                                                 \
4334    SLANG_FORCE_INLINE SLANG_CUDA_CALL T##2 tex1Dfetch_int(CUtexObject texObj, int x, int mip) \
4335    {                                                                                          \
4336        T result_x, result_y;                                                                  \
4337        T stub;                                                                                \
4338        asm("tex.level.1d.v4." dtype ".s32 {%0, %1, %2, %3}, [%4, {%5}], %6;"                  \
4339            : c(result_x), c(result_y), c(stub), c(stub)                                       \
4340            : "l"(texObj), "r"(x), "r"(mip));                                                  \
4341        return make_##T##2(result_x, result_y);                                                \
4342    }                                                                                          \
4343    template<>                                                                                 \
4344    SLANG_FORCE_INLINE SLANG_CUDA_CALL T##4 tex1Dfetch_int(CUtexObject texObj, int x, int mip) \
4345    {                                                                                          \
4346        T result_x, result_y, result_z, result_w;                                              \
4347        asm("tex.level.1d.v4." dtype ".s32 {%0, %1, %2, %3}, [%4, {%5}], %6;"                  \
4348            : c(result_x), c(result_y), c(result_z), c(result_w)                               \
4349            : "l"(texObj), "r"(x), "r"(mip));                                                  \
4350        return make_##T##4(result_x, result_y, result_z, result_w);                            \
4351    }
4352
4353SLANG_TEX1DFETCH_INT_IMPL(float, "f32", "=f")
4354SLANG_TEX1DFETCH_INT_IMPL(uint, "u32", "=r")
4355SLANG_TEX1DFETCH_INT_IMPL(int, "s32", "=r")
4356#endif
4357
4358template<typename T>
4359SLANG_FORCE_INLINE SLANG_CUDA_CALL T tex2Dfetch_int(CUtexObject texObj, int x, int y, int mip);
4360
4361#define SLANG_TEX2DFETCH_INT_IMPL(T, dtype, c)                                                     \
4362    template<>                                                                                     \
4363    SLANG_FORCE_INLINE SLANG_CUDA_CALL T tex2Dfetch_int(CUtexObject texObj, int x, int y, int mip) \
4364    {                                                                                              \
4365        T result;                                                                                  \
4366        T stub;                                                                                    \
4367        asm("tex.level.2d.v4." dtype ".s32 {%0, %1, %2, %3}, [%4, {%5, %6}], %7;"                  \
4368            : c(result), c(stub), c(stub), c(stub)                                                 \
4369            : "l"(texObj), "r"(x), "r"(y), "r"(mip));                                              \
4370        return result;                                                                             \
4371    }                                                                                              \
4372    template<>                                                                                     \
4373    SLANG_FORCE_INLINE SLANG_CUDA_CALL                                                             \
4374        T##2 tex2Dfetch_int(CUtexObject texObj, int x, int y, int mip)                             \
4375    {                                                                                              \
4376        T result_x, result_y;                                                                      \
4377        T stub;                                                                                    \
4378        asm("tex.level.2d.v4." dtype ".s32 {%0, %1, %2, %3}, [%4, {%5, %6}], %7;"                  \
4379            : c(result_x), c(result_y), c(stub), c(stub)                                           \
4380            : "l"(texObj), "r"(x), "r"(y), "r"(mip));                                              \
4381        return make_##T##2(result_x, result_y);                                                    \
4382    }                                                                                              \
4383    template<>                                                                                     \
4384    SLANG_FORCE_INLINE SLANG_CUDA_CALL                                                             \
4385        T##4 tex2Dfetch_int(CUtexObject texObj, int x, int y, int mip)                             \
4386    {                                                                                              \
4387        T result_x, result_y, result_z, result_w;                                                  \
4388        asm("tex.level.2d.v4." dtype ".s32 {%0, %1, %2, %3}, [%4, {%5, %6}], %7;"                  \
4389            : c(result_x), c(result_y), c(result_z), c(result_w)                                   \
4390            : "l"(texObj), "r"(x), "r"(y), "r"(mip));                                              \
4391        return make_##T##4(result_x, result_y, result_z, result_w);                                \
4392    }
4393
4394SLANG_TEX2DFETCH_INT_IMPL(float, "f32", "=f")
4395SLANG_TEX2DFETCH_INT_IMPL(uint, "u32", "=r")
4396SLANG_TEX2DFETCH_INT_IMPL(int, "s32", "=r")
4397
4398
4399template<typename T>
4400SLANG_FORCE_INLINE SLANG_CUDA_CALL T
4401tex3Dfetch_int(CUtexObject texObj, int x, int y, int z, int mip);
4402
4403#define SLANG_TEX3DFETCH_INT_IMPL(T, dtype, c)                                            \
4404    template<>                                                                            \
4405    SLANG_FORCE_INLINE SLANG_CUDA_CALL T                                                  \
4406    tex3Dfetch_int(CUtexObject texObj, int x, int y, int z, int mip)                      \
4407    {                                                                                     \
4408        T result;                                                                         \
4409        T stub;                                                                           \
4410        asm("tex.level.3d.v4." dtype ".s32 {%0, %1, %2, %3}, [%4, {%5, %6, %7, %8}], %9;" \
4411            : c(result), c(stub), c(stub), c(stub)                                        \
4412            : "l"(texObj), "r"(x), "r"(y), "r"(z), "r"(z) /* ignored */, "r"(mip));       \
4413        return result;                                                                    \
4414    }                                                                                     \
4415    template<>                                                                            \
4416    SLANG_FORCE_INLINE SLANG_CUDA_CALL                                                    \
4417        T##2 tex3Dfetch_int(CUtexObject texObj, int x, int y, int z, int mip)             \
4418    {                                                                                     \
4419        T result_x, result_y;                                                             \
4420        T stub;                                                                           \
4421        asm("tex.level.3d.v4." dtype ".s32 {%0, %1, %2, %3}, [%4, {%5, %6, %7, %8}], %9;" \
4422            : c(result_x), c(result_y), c(stub), c(stub)                                  \
4423            : "l"(texObj), "r"(x), "r"(y), "r"(z), "r"(z) /* ignored */, "r"(mip));       \
4424        return make_##T##2(result_x, result_y);                                           \
4425    }                                                                                     \
4426    template<>                                                                            \
4427    SLANG_FORCE_INLINE SLANG_CUDA_CALL                                                    \
4428        T##4 tex3Dfetch_int(CUtexObject texObj, int x, int y, int z, int mip)             \
4429    {                                                                                     \
4430        T result_x, result_y, result_z, result_w;                                         \
4431        asm("tex.level.3d.v4." dtype ".s32 {%0, %1, %2, %3}, [%4, {%5, %6, %7, %8}], %9;" \
4432            : c(result_x), c(result_y), c(result_z), c(result_w)                          \
4433            : "l"(texObj), "r"(x), "r"(y), "r"(z), "r"(z) /* ignored */, "r"(mip));       \
4434        return make_##T##4(result_x, result_y, result_z, result_w);                       \
4435    }
4436
4437SLANG_TEX3DFETCH_INT_IMPL(float, "f32", "=f")
4438SLANG_TEX3DFETCH_INT_IMPL(uint, "u32", "=r")
4439SLANG_TEX3DFETCH_INT_IMPL(int, "s32", "=r")
4440
4441template<typename T>
4442SLANG_FORCE_INLINE SLANG_CUDA_CALL T
4443tex1DArrayfetch_int(CUtexObject texObj, int x, int layer, int mip);
4444
4445#define SLANG_TEX1DARRAYFETCH_INT_IMPL(T, dtype, c)                                \
4446    template<>                                                                     \
4447    SLANG_FORCE_INLINE SLANG_CUDA_CALL T                                           \
4448    tex1DArrayfetch_int(CUtexObject texObj, int x, int layer, int mip)             \
4449    {                                                                              \
4450        T result;                                                                  \
4451        T stub;                                                                    \
4452        asm("tex.level.a1d.v4." dtype ".s32 {%0, %1, %2, %3}, [%4, {%5, %6}], %7;" \
4453            : c(result), c(stub), c(stub), c(stub)                                 \
4454            : "l"(texObj), "r"(layer), "r"(x), "r"(mip));                          \
4455        return result;                                                             \
4456    }                                                                              \
4457    template<>                                                                     \
4458    SLANG_FORCE_INLINE SLANG_CUDA_CALL                                             \
4459        T##2 tex1DArrayfetch_int(CUtexObject texObj, int x, int layer, int mip)    \
4460    {                                                                              \
4461        T result_x, result_y;                                                      \
4462        T stub;                                                                    \
4463        asm("tex.level.a1d.v4." dtype ".s32 {%0, %1, %2, %3}, [%4, {%5, %6}], %7;" \
4464            : c(result_x), c(result_y), c(stub), c(stub)                           \
4465            : "l"(texObj), "r"(layer), "r"(x), "r"(mip));                          \
4466        return make_##T##2(result_x, result_y);                                    \
4467    }                                                                              \
4468    template<>                                                                     \
4469    SLANG_FORCE_INLINE SLANG_CUDA_CALL                                             \
4470        T##4 tex1DArrayfetch_int(CUtexObject texObj, int x, int layer, int mip)    \
4471    {                                                                              \
4472        T result_x, result_y, result_z, result_w;                                  \
4473        asm("tex.level.a1d.v4." dtype ".s32 {%0, %1, %2, %3}, [%4, {%5, %6}], %7;" \
4474            : c(result_x), c(result_y), c(result_z), c(result_w)                   \
4475            : "l"(texObj), "r"(layer), "r"(x), "r"(mip));                          \
4476        return make_##T##4(result_x, result_y, result_z, result_w);                \
4477    }
4478
4479SLANG_TEX1DARRAYFETCH_INT_IMPL(float, "f32", "=f")
4480SLANG_TEX1DARRAYFETCH_INT_IMPL(uint, "u32", "=r")
4481SLANG_TEX1DARRAYFETCH_INT_IMPL(int, "s32", "=r")
4482
4483template<typename T>
4484SLANG_FORCE_INLINE SLANG_CUDA_CALL T
4485tex2DArrayfetch_int(CUtexObject texObj, int x, int y, int layer, int mip);
4486
4487#define SLANG_TEX2DARRAYFETCH_INT_IMPL(T, dtype, c)                                         \
4488    template<>                                                                              \
4489    SLANG_FORCE_INLINE SLANG_CUDA_CALL T                                                    \
4490    tex2DArrayfetch_int(CUtexObject texObj, int x, int y, int layer, int mip)               \
4491    {                                                                                       \
4492        T result;                                                                           \
4493        T stub;                                                                             \
4494        asm("tex.level.a2d.v4." dtype ".s32 {%0, %1, %2, %3}, [%4, {%5, %6, %7, %8}], %9;"  \
4495            : c(result), c(stub), c(stub), c(stub)                                          \
4496            : "l"(texObj), "r"(layer), "r"(x), "r"(y), "r"(layer) /* ignored */, "r"(mip)); \
4497        return result;                                                                      \
4498    }                                                                                       \
4499    template<>                                                                              \
4500    SLANG_FORCE_INLINE SLANG_CUDA_CALL                                                      \
4501        T##2 tex2DArrayfetch_int(CUtexObject texObj, int x, int y, int layer, int mip)      \
4502    {                                                                                       \
4503        T result_x, result_y;                                                               \
4504        T stub;                                                                             \
4505        asm("tex.level.a2d.v4." dtype ".s32 {%0, %1, %2, %3}, [%4, {%5, %6, %7, %8}], %9;"  \
4506            : c(result_x), c(result_y), c(stub), c(stub)                                    \
4507            : "l"(texObj), "r"(layer), "r"(x), "r"(y), "r"(layer) /* ignored */, "r"(mip)); \
4508        return make_##T##2(result_x, result_y);                                             \
4509    }                                                                                       \
4510    template<>                                                                              \
4511    SLANG_FORCE_INLINE SLANG_CUDA_CALL                                                      \
4512        T##4 tex2DArrayfetch_int(CUtexObject texObj, int x, int y, int layer, int mip)      \
4513    {                                                                                       \
4514        T result_x, result_y, result_z, result_w;                                           \
4515        asm("tex.level.a2d.v4." dtype ".s32 {%0, %1, %2, %3}, [%4, {%5, %6, %7, %8}], %9;"  \
4516            : c(result_x), c(result_y), c(result_z), c(result_w)                            \
4517            : "l"(texObj), "r"(layer), "r"(x), "r"(y), "r"(layer) /* ignored */, "r"(mip)); \
4518        return make_##T##4(result_x, result_y, result_z, result_w);                         \
4519    }
4520
4521SLANG_TEX2DARRAYFETCH_INT_IMPL(float, "f32", "=f")
4522SLANG_TEX2DARRAYFETCH_INT_IMPL(uint, "u32", "=r")
4523SLANG_TEX2DARRAYFETCH_INT_IMPL(int, "s32", "=r")
4524
4525// Wave rotate helper functions - templated approach
4526#define SLANG_WARP_FULL_MASK 0xFFFFFFFF
4527
4528// Macro-based wave rotate implementation following codebase patterns
4529#define SLANG_WAVE_ROTATE_IMPL(T)                                                     \
4530    __device__ __forceinline__ T##2 _slang_waveRotate(T##2 value, unsigned int delta) \
4531    {                                                                                 \
4532        return make_##T##2(                                                           \
4533            (T)__shfl_sync(                                                           \
4534                SLANG_WARP_FULL_MASK,                                                 \
4535                value.x,                                                              \
4536                (_getLaneId() + delta) % SLANG_CUDA_WARP_SIZE),                       \
4537            (T)__shfl_sync(                                                           \
4538                SLANG_WARP_FULL_MASK,                                                 \
4539                value.y,                                                              \
4540                (_getLaneId() + delta) % SLANG_CUDA_WARP_SIZE));                      \
4541    }                                                                                 \
4542    __device__ __forceinline__ T##3 _slang_waveRotate(T##3 value, unsigned int delta) \
4543    {                                                                                 \
4544        return make_##T##3(                                                           \
4545            (T)__shfl_sync(                                                           \
4546                SLANG_WARP_FULL_MASK,                                                 \
4547                value.x,                                                              \
4548                (_getLaneId() + delta) % SLANG_CUDA_WARP_SIZE),                       \
4549            (T)__shfl_sync(                                                           \
4550                SLANG_WARP_FULL_MASK,                                                 \
4551                value.y,                                                              \
4552                (_getLaneId() + delta) % SLANG_CUDA_WARP_SIZE),                       \
4553            (T)__shfl_sync(                                                           \
4554                SLANG_WARP_FULL_MASK,                                                 \
4555                value.z,                                                              \
4556                (_getLaneId() + delta) % SLANG_CUDA_WARP_SIZE));                      \
4557    }                                                                                 \
4558    __device__ __forceinline__ T##4 _slang_waveRotate(T##4 value, unsigned int delta) \
4559    {                                                                                 \
4560        return make_##T##4(                                                           \
4561            (T)__shfl_sync(                                                           \
4562                SLANG_WARP_FULL_MASK,                                                 \
4563                value.x,                                                              \
4564                (_getLaneId() + delta) % SLANG_CUDA_WARP_SIZE),                       \
4565            (T)__shfl_sync(                                                           \
4566                SLANG_WARP_FULL_MASK,                                                 \
4567                value.y,                                                              \
4568                (_getLaneId() + delta) % SLANG_CUDA_WARP_SIZE),                       \
4569            (T)__shfl_sync(                                                           \
4570                SLANG_WARP_FULL_MASK,                                                 \
4571                value.z,                                                              \
4572                (_getLaneId() + delta) % SLANG_CUDA_WARP_SIZE),                       \
4573            (T)__shfl_sync(                                                           \
4574                SLANG_WARP_FULL_MASK,                                                 \
4575                value.w,                                                              \
4576                (_getLaneId() + delta) % SLANG_CUDA_WARP_SIZE));                      \
4577    }
4578
4579// Generate wave rotate functions for all standard vector types
4580SLANG_WAVE_ROTATE_IMPL(uint)
4581SLANG_WAVE_ROTATE_IMPL(int)
4582SLANG_WAVE_ROTATE_IMPL(float)
4583SLANG_WAVE_ROTATE_IMPL(short)
4584SLANG_WAVE_ROTATE_IMPL(ushort)
4585SLANG_WAVE_ROTATE_IMPL(char)
4586SLANG_WAVE_ROTATE_IMPL(uchar)
4587SLANG_WAVE_ROTATE_IMPL(longlong)
4588SLANG_WAVE_ROTATE_IMPL(ulonglong)
4589
4590#ifdef SLANG_CUDA_ENABLE_HALF
4591SLANG_WAVE_ROTATE_IMPL(__half)
4592#endif
4593
4594// Special handling for boolean vectors (requires int conversion)
4595__device__ __forceinline__ bool2 _slang_waveRotate(bool2 value, unsigned int delta)
4596{
4597    int2 intValue = make_int2((int)value.x, (int)value.y);
4598    int2 result = _slang_waveRotate(intValue, delta);
4599    return make_bool2((bool)result.x, (bool)result.y);
4600}
4601
4602__device__ __forceinline__ bool3 _slang_waveRotate(bool3 value, unsigned int delta)
4603{
4604    int3 intValue = make_int3((int)value.x, (int)value.y, (int)value.z);
4605    int3 result = _slang_waveRotate(intValue, delta);
4606    return make_bool3((bool)result.x, (bool)result.y, (bool)result.z);
4607}
4608
4609__device__ __forceinline__ bool4 _slang_waveRotate(bool4 value, unsigned int delta)
4610{
4611    int4 intValue = make_int4((int)value.x, (int)value.y, (int)value.z, (int)value.w);
4612    int4 result = _slang_waveRotate(intValue, delta);
4613    return make_bool4((bool)result.x, (bool)result.y, (bool)result.z, (bool)result.w);
4614}
4615
4616#undef SLANG_WAVE_ROTATE_IMPL
4617
4618// Quad control operations for CUDA
4619__device__ __forceinline__ bool _slang_quadAny(bool expr)
4620{
4621    // Get values from all 4 lanes in the quad
4622    bool v0 = __shfl_sync(0xFFFFFFFF, expr, (_getLaneId() & 0xFFFFFFFC) | 0);
4623    bool v1 = __shfl_sync(0xFFFFFFFF, expr, (_getLaneId() & 0xFFFFFFFC) | 1);
4624    bool v2 = __shfl_sync(0xFFFFFFFF, expr, (_getLaneId() & 0xFFFFFFFC) | 2);
4625    bool v3 = __shfl_sync(0xFFFFFFFF, expr, (_getLaneId() & 0xFFFFFFFC) | 3);
4626    return v0 || v1 || v2 || v3;
4627}
4628
4629__device__ __forceinline__ bool _slang_quadAll(bool expr)
4630{
4631    // Get values from all 4 lanes in the quad
4632    bool v0 = __shfl_sync(0xFFFFFFFF, expr, (_getLaneId() & 0xFFFFFFFC) | 0);
4633    bool v1 = __shfl_sync(0xFFFFFFFF, expr, (_getLaneId() & 0xFFFFFFFC) | 1);
4634    bool v2 = __shfl_sync(0xFFFFFFFF, expr, (_getLaneId() & 0xFFFFFFFC) | 2);
4635    bool v3 = __shfl_sync(0xFFFFFFFF, expr, (_getLaneId() & 0xFFFFFFFC) | 3);
4636    return v0 && v1 && v2 && v3;
4637}
4638
4639// Clustered wave rotate operations for CUDA
4640// Clustered rotate rotates values within clusters of specified size
4641#define SLANG_WAVE_CLUSTERED_ROTATE_IMPL(T)                                                       \
4642    __device__ __forceinline__ T                                                                  \
4643    _slang_waveClusteredRotate(T value, unsigned int delta, unsigned int clusterSize)             \
4644    {                                                                                             \
4645        unsigned int laneId = _getLaneId();                                                       \
4646        unsigned int clusterStart = (laneId / clusterSize) * clusterSize;                         \
4647        unsigned int targetLane = clusterStart + ((laneId - clusterStart + delta) % clusterSize); \
4648        return __shfl_sync(SLANG_WARP_FULL_MASK, value, targetLane);                              \
4649    }                                                                                             \
4650    __device__ __forceinline__                                                                    \
4651        T##2 _slang_waveClusteredRotate(T##2 value, unsigned int delta, unsigned int clusterSize) \
4652    {                                                                                             \
4653        unsigned int laneId = _getLaneId();                                                       \
4654        unsigned int clusterStart = (laneId / clusterSize) * clusterSize;                         \
4655        unsigned int targetLane = clusterStart + ((laneId - clusterStart + delta) % clusterSize); \
4656        return make_##T##2(                                                                       \
4657            (T)__shfl_sync(SLANG_WARP_FULL_MASK, value.x, targetLane),                            \
4658            (T)__shfl_sync(SLANG_WARP_FULL_MASK, value.y, targetLane));                           \
4659    }                                                                                             \
4660    __device__ __forceinline__                                                                    \
4661        T##3 _slang_waveClusteredRotate(T##3 value, unsigned int delta, unsigned int clusterSize) \
4662    {                                                                                             \
4663        unsigned int laneId = _getLaneId();                                                       \
4664        unsigned int clusterStart = (laneId / clusterSize) * clusterSize;                         \
4665        unsigned int targetLane = clusterStart + ((laneId - clusterStart + delta) % clusterSize); \
4666        return make_##T##3(                                                                       \
4667            (T)__shfl_sync(SLANG_WARP_FULL_MASK, value.x, targetLane),                            \
4668            (T)__shfl_sync(SLANG_WARP_FULL_MASK, value.y, targetLane),                            \
4669            (T)__shfl_sync(SLANG_WARP_FULL_MASK, value.z, targetLane));                           \
4670    }                                                                                             \
4671    __device__ __forceinline__                                                                    \
4672        T##4 _slang_waveClusteredRotate(T##4 value, unsigned int delta, unsigned int clusterSize) \
4673    {                                                                                             \
4674        unsigned int laneId = _getLaneId();                                                       \
4675        unsigned int clusterStart = (laneId / clusterSize) * clusterSize;                         \
4676        unsigned int targetLane = clusterStart + ((laneId - clusterStart + delta) % clusterSize); \
4677        return make_##T##4(                                                                       \
4678            (T)__shfl_sync(SLANG_WARP_FULL_MASK, value.x, targetLane),                            \
4679            (T)__shfl_sync(SLANG_WARP_FULL_MASK, value.y, targetLane),                            \
4680            (T)__shfl_sync(SLANG_WARP_FULL_MASK, value.z, targetLane),                            \
4681            (T)__shfl_sync(SLANG_WARP_FULL_MASK, value.w, targetLane));                           \
4682    }
4683
4684// Generate clustered wave rotate functions for all standard types
4685SLANG_WAVE_CLUSTERED_ROTATE_IMPL(uint)
4686SLANG_WAVE_CLUSTERED_ROTATE_IMPL(int)
4687SLANG_WAVE_CLUSTERED_ROTATE_IMPL(float)
4688SLANG_WAVE_CLUSTERED_ROTATE_IMPL(short)
4689SLANG_WAVE_CLUSTERED_ROTATE_IMPL(ushort)
4690SLANG_WAVE_CLUSTERED_ROTATE_IMPL(char)
4691SLANG_WAVE_CLUSTERED_ROTATE_IMPL(uchar)
4692SLANG_WAVE_CLUSTERED_ROTATE_IMPL(longlong)
4693SLANG_WAVE_CLUSTERED_ROTATE_IMPL(ulonglong)
4694
4695#ifdef SLANG_CUDA_ENABLE_HALF
4696SLANG_WAVE_CLUSTERED_ROTATE_IMPL(__half)
4697#endif
4698
4699// Special handling for boolean clustered rotate
4700__device__ __forceinline__ bool _slang_waveClusteredRotate(
4701    bool value,
4702    unsigned int delta,
4703    unsigned int clusterSize)
4704{
4705    int intValue = (int)value;
4706    int result = _slang_waveClusteredRotate(intValue, delta, clusterSize);
4707    return (bool)result;
4708}
4709
4710__device__ __forceinline__ bool2
4711_slang_waveClusteredRotate(bool2 value, unsigned int delta, unsigned int clusterSize)
4712{
4713    int2 intValue = make_int2((int)value.x, (int)value.y);
4714    int2 result = _slang_waveClusteredRotate(intValue, delta, clusterSize);
4715    return make_bool2((bool)result.x, (bool)result.y);
4716}
4717
4718__device__ __forceinline__ bool3
4719_slang_waveClusteredRotate(bool3 value, unsigned int delta, unsigned int clusterSize)
4720{
4721    int3 intValue = make_int3((int)value.x, (int)value.y, (int)value.z);
4722    int3 result = _slang_waveClusteredRotate(intValue, delta, clusterSize);
4723    return make_bool3((bool)result.x, (bool)result.y, (bool)result.z);
4724}
4725
4726__device__ __forceinline__ bool4
4727_slang_waveClusteredRotate(bool4 value, unsigned int delta, unsigned int clusterSize)
4728{
4729    int4 intValue = make_int4((int)value.x, (int)value.y, (int)value.z, (int)value.w);
4730    int4 result = _slang_waveClusteredRotate(intValue, delta, clusterSize);
4731    return make_bool4((bool)result.x, (bool)result.y, (bool)result.z, (bool)result.w);
4732}
4733
4734#undef SLANG_WAVE_CLUSTERED_ROTATE_IMPL