src: gpu: intel: include: make philox overloadable

Author: h-sadia

src/gpu/intel/eltwise/ref.cl

@@ -99,10 +99,10 @@ __kernel void ref_eltwise_fwd(__global SRC_DATA_T *src,
     uint dropout_threshold = get_dropout_threshold(dropout_p);
     float dropout_inv_q = (dropout_p != 1.f) ? 1.f / (1.f - dropout_p) : 0.f;
 #if USE_OFFSET
-    uint res = philox_4x32_u64_w_offset(
+    uint res = philox_4x32_w_offset(
             (ulong)data_off, (ulong)dropout_seed, (ulong)dropout_offset);
 #else
-    uint res = philox_4x32_u64((ulong)data_off, (ulong)dropout_seed);
+    uint res = philox_4x32((ulong)data_off, (ulong)dropout_seed);
 #endif
     uchar dropout = res > dropout_threshold;
     tmp_s = (dropout) ? tmp_s * dropout_inv_q : 0;

src/gpu/intel/gemm/with_post_ops.cl

@@ -123,10 +123,10 @@ __kernel void gemm_post_ops(__global SRC_DATA_T *src,
     uint dropout_threshold = get_dropout_threshold(dropout_p);
     float dropout_inv_q = (dropout_p != 1.f) ? 1.f / (1.f - dropout_p) : 0.f;
 #if DROPOUT_USE_OFFSET
-    uint res = philox_4x32_u64_w_offset(
+    uint res = philox_4x32_w_offset(
             (ulong)data_idx, (ulong)dropout_seed, (ulong)dropout_offset);
 #else
-    uint res = philox_4x32_u64((ulong)data_idx, (ulong)dropout_seed);
+    uint res = philox_4x32((ulong)data_idx, (ulong)dropout_seed);
 #endif
     uchar dropout = res > dropout_threshold;
     accumulator = (dropout) ? accumulator * dropout_inv_q : 0;

src/gpu/intel/include/philox.h

@@ -20,7 +20,7 @@
 #define DT_UNDEF 1
 #include "gpu/intel/include/types.h"
 
-uint4 philox_4x32_s64_vec4(ulong idx, ulong seed, ulong offset) {
+uint4 philox_4x32_vec4_w_offset(ulong idx, ulong seed, ulong offset) {
 #define PHILOX_4UINT_ROUND(mul, ctr, key) \
     as_uint4(convert_ulong2(ctr.s02) * mul).s3210 \
             ^ (uint4)(ctr.s1 ^ key.s0, 0, ctr.s3 ^ key.s1, 0)
@@ -53,40 +53,45 @@ uint4 philox_4x32_s64_vec4(ulong idx, ulong seed, ulong offset) {
     return ctr;
 }
 
-uint philox_4x32_u64_w_offset(ulong idx, ulong seed, ulong offset) {
-    return philox_4x32_s64_vec4(idx, seed, offset)[idx & 3L];
+uint philox_4x32_w_offset(ulong idx, ulong seed, ulong offset) {
+    return philox_4x32_vec4_w_offset(idx, seed, offset)[idx & 3L];
 }
 
-uint philox_4x32_u64(ulong idx, ulong seed) {
+uint __attribute__((overloadable)) philox_4x32(ulong idx, ulong seed) {
     // Note: this is for compatibility with impls that don't support s64 rand
     ulong x = idx & ~3L;
     ulong idx_64 = ((x + 3) << 32) + (x + 2);
     ulong offset_64 = ((x + 1) << 32) + x;
     ulong seed_64 = (seed << 32) + seed;
-    return philox_4x32_s64(idx_64, seed_64, offset_64);
+    return philox_4x32_w_offset(idx_64, seed_64, offset_64);
 }
 
-uint philox_4x32(uint idx, uint seed) {
-    // Note: preserve old signature for compatibility
-    return philox_4x32_u64((ulong)idx, (ulong)seed);
+uint __attribute__((overloadable)) philox_4x32(long idx, long seed) {
+    // Convert long to ulong and call the existing function
+    return philox_4x32((ulong)idx, (ulong)seed);
+}
+
+uint __attribute__((overloadable)) philox_4x32(int idx, int seed) {
+    // Convert int to ulong and call the existing overloadable function
+    return philox_4x32((ulong)idx, (ulong)seed);
 }
 
 uint4 philox_4x32_vec4(uint idx, uint seed) {
     ulong x = idx & ~3L;
     ulong idx_64 = ((x + 3) << 32) + (x + 2);
     ulong offset_64 = ((x + 1) << 32) + x;
     ulong seed_64 = ((ulong)(seed) << 32) + seed;
-    return philox_4x32_s64_vec4(idx_64, seed_64, offset_64);
+    return philox_4x32_vec4_w_offset(idx_64, seed_64, offset_64);
 }
 
 ushort philox_8x16(long idx, uint seed) {
     ulong idx_ = (ulong)idx;
-    return as_ushort2(philox_4x32_u64(idx_ >> 1, (ulong)seed))[idx_ & 1];
+    return as_ushort2(philox_4x32(idx_ >> 1, (ulong)seed))[idx_ & 1];
 }
 
 uchar philox_16x8(long idx, uint seed) {
     ulong idx_ = (ulong)idx;
-    return as_uchar4(philox_4x32_u64(idx_ >> 2, (ulong)seed))[idx_ & 3];
+    return as_uchar4(philox_4x32(idx_ >> 2, (ulong)seed))[idx_ & 3];
 }
 
 #if WITH_SROUND
@@ -120,4 +125,7 @@ uint get_dropout_threshold(float p) {
             + !!(mantissa & ((1u << exponent) - 1u));
 }
 #endif
+
+
+
 #endif

src/gpu/intel/matmul/ref.cl

@@ -289,10 +289,10 @@ __kernel void ref_matmul(__global SRC_DATA_T *A, __global WEI_DATA_T *B,
 
 #if WITH_DROPOUT
 #if WITH_SEED_S64 && USE_OFFSET
-    uint res = philox_4x32_u64_w_offset(
+    uint res = philox_4x32_w_offset(
             (ulong)dst_off, (ulong)dropout_seed, (ulong)dropout_offset);
 #else
-    uint res = philox_4x32_u64((ulong)dst_off, (ulong)dropout_seed);
+    uint res = philox_4x32((ulong)dst_off, (ulong)dropout_seed);
 #endif
     uchar dropout = res > dropout_threshold;
     po_acc = (dropout) ? po_acc * dropout_inv_q : 0;

src/gpu/intel/sdpa/micro.cl

@@ -56,9 +56,9 @@ inline void apply_dropout_s_tile(
         ulong _goff = batch_head_base + (ulong)offset_c * (ulong)k_stride \
                 + (ulong)offset_r; \
         uint _philox = use_dropout_offset \
-                ? philox_4x32_u64_w_offset( \
+                ? philox_4x32_w_offset( \
                           (ulong)_goff, (ulong)seed, (ulong)offset) \
-                : philox_4x32_u64((ulong)_goff, (ulong)seed); \
+                : philox_4x32((ulong)_goff, (ulong)seed); \
         (offset_r < max_r && offset_c < max_c) && (_philox > threshold); \
     })
 

src/gpu/intel/sdpa/micro_bwd.cl

@@ -52,9 +52,9 @@ typedef ugemm_ktq_c_type ktq_tile_type; // D*Br tile
         ulong _goff = batch_head_base + (ulong)offset_c * (ulong)k_stride \
                 + (ulong)offset_r; \
         uint _philox = use_dropout_offset \
-                ? philox_4x32_u64_w_offset( \
+                ? philox_4x32_w_offset( \
                           (ulong)_goff, (ulong)seed, (ulong)offset) \
-                : philox_4x32_u64((ulong)_goff, (ulong)seed); \
+                : philox_4x32((ulong)_goff, (ulong)seed); \
         (offset_r < max_r && offset_c < max_c) && (_philox > threshold); \
     })
 

src/gpu/intel/softmax/simple.cl

@@ -190,10 +190,10 @@ simple_softmax_fwd_generic(__global SRC_DATA_T *src, __global DATA_T *dst,
         float dropout_inv_q
                 = (dropout_p != 1.f) ? 1.f / (1.f - dropout_p) : 0.f;
 #if USE_OFFSET
-        uint res = philox_4x32_u64_w_offset(
+        uint res = philox_4x32_w_offset(
                 (ulong)data_off, (ulong)dropout_seed, (ulong)dropout_offset);
 #else
-        uint res = philox_4x32_u64((ulong)data_off, (ulong)dropout_seed);
+        uint res = philox_4x32((ulong)data_off, (ulong)dropout_seed);
 #endif
         uchar dropout = res > dropout_threshold;
         tmp = (dropout) ? tmp * dropout_inv_q : 0;