Use uint3 for both fastdiv and fastmodulo

The compiler seems to reliably optimize away the unused .z component in
the fastdiv use-case, see https://godbolt.org/z/rx8KPrKr3

Author: ORippler

ggml/src/ggml-cuda/common.cuh

@@ -569,33 +569,30 @@ static __device__ __forceinline__ float ggml_cuda_e8m0_to_fp32(uint8_t x) {
 // and a shift:
 //
 // n/d = (mulhi(n, mp) + n) >> L;
-static const uint2 init_fastdiv_values(uint32_t d) {
+static const uint3 init_fastdiv_values(uint32_t d) {
     // compute L = ceil(log2(d));
     uint32_t L = 0;
     while (L < 32 && (uint32_t{ 1 } << L) < d) {
         L++;
     }
 
     uint32_t mp = (uint32_t) ((uint64_t{ 1 } << 32) * ((uint64_t{ 1 } << L) - d) / d + 1);
-    return make_uint2(mp, L);
+    // pack divisor as well to reduce error surface
+    return make_uint3(mp, L, d);
 }
 
-static __device__ __forceinline__ uint32_t fastdiv(uint32_t n, const uint2 div_consts) {
+static __device__ __forceinline__ uint32_t fastdiv(uint32_t n, const uint3 div_consts) {
+    // expects div_consts to contain <mp, L, divisor> in <x, y, z>
+    // div_consts.z is unused and optimized away by the compiler.
     // Compute high 32 bits of n * mp
     const uint32_t hi = __umulhi(n, div_consts.x);
-    // Apply the formula
+    // add n, apply bit shift
     return (hi + n) >> div_consts.y;
 }
 
-static const uint3 init_fastmodulo_values(uint32_t d) {
-    // uint3 contains <mp, L, divisor> in <x, y, z>
-    const uint2 fastdiv = init_fastdiv_values(d);
-    return make_uint3(fastdiv.x, fastdiv.y, d);
-}
-
 static __device__ __forceinline__ uint32_t fastmodulo(uint32_t n, const uint3 modulo_consts) {
-    // expects modulo_consts to contain <mp, L, divisor> in <x, y, z> (see init_fastmodulo_values function)
-    return n - fastdiv(n, make_uint2(modulo_consts.x, modulo_consts.y)) * modulo_consts.z;
+    // expects modulo_consts to contain <mp, L, divisor> in <x, y, z> (see init_fastdiv_values)
+    return n - fastdiv(n, modulo_consts) * modulo_consts.z;
 }
 
 typedef void (*dequantize_kernel_t)(const void * vx, const int64_t ib, const int iqs, float2 & v);

ggml/src/ggml-cuda/norm.cu

@@ -398,10 +398,10 @@ static void rms_norm_mul_f32_cuda(const float *  x,
         return;
     }
     if (add == nullptr) {
-        uint3 mul_ncols_packed     = init_fastmodulo_values(mul_ncols);
-        uint3 mul_nrows_packed     = init_fastmodulo_values(mul_nrows);
-        uint3 mul_nchannels_packed = init_fastmodulo_values(mul_nchannels);
-        uint3 mul_nsamples_packed  = init_fastmodulo_values(mul_nsamples);
+        uint3 mul_ncols_packed     = init_fastdiv_values(mul_ncols);
+        uint3 mul_nrows_packed     = init_fastdiv_values(mul_nrows);
+        uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels);
+        uint3 mul_nsamples_packed  = init_fastdiv_values(mul_nsamples);
         if (ncols < 1024) {
             const dim3 block_dims(256, 1, 1);
             rms_norm_f32<256, true><<<blocks_num, block_dims, 0, stream>>>(
@@ -414,15 +414,15 @@ static void rms_norm_mul_f32_cuda(const float *  x,
                 mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed);
         }
     } else {
-        uint3 mul_ncols_packed     = init_fastmodulo_values(mul_ncols);
-        uint3 mul_nrows_packed     = init_fastmodulo_values(mul_nrows);
-        uint3 mul_nchannels_packed = init_fastmodulo_values(mul_nchannels);
-        uint3 mul_nsamples_packed  = init_fastmodulo_values(mul_nsamples);
-
-        uint3 add_ncols_packed     = init_fastmodulo_values(add_ncols);
-        uint3 add_nrows_packed     = init_fastmodulo_values(add_nrows);
-        uint3 add_nchannels_packed = init_fastmodulo_values(add_nchannels);
-        uint3 add_nsamples_packed  = init_fastmodulo_values(add_nsamples);
+        uint3 mul_ncols_packed     = init_fastdiv_values(mul_ncols);
+        uint3 mul_nrows_packed     = init_fastdiv_values(mul_nrows);
+        uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels);
+        uint3 mul_nsamples_packed  = init_fastdiv_values(mul_nsamples);
+
+        uint3 add_ncols_packed     = init_fastdiv_values(add_ncols);
+        uint3 add_nrows_packed     = init_fastdiv_values(add_nrows);
+        uint3 add_nchannels_packed = init_fastdiv_values(add_nchannels);
+        uint3 add_nsamples_packed  = init_fastdiv_values(add_nsamples);
         if (ncols < 1024) {
             const dim3 block_dims(256, 1, 1);
             rms_norm_f32<256, true, true><<<blocks_num, block_dims, 0, stream>>>(