faster-diffbloch 0.1.0__py3-none-any.whl

faster_diffbloch/native/batch_cgemm.metal

@@ -0,0 +1,374 @@
+#include <metal_stdlib>
+using namespace metal;
+
+static inline float2 c_mul(float2 a, float2 b) {
+    return float2(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x);
+}
+
+static inline float2 c_fma(float2 a, float2 b, float2 acc) {
+    return float2(acc.x + (a.x * b.x - a.y * b.y), acc.y + (a.x * b.y + a.y * b.x));
+}
+
+static inline float2 c_conj(float2 a) {
+    return float2(a.x, -a.y);
+}
+
+constant int TILE_SIZE = 16;
+
+struct BatchGemmParams {
+    int M;
+    int N;
+    int K;
+    int lda;
+    int ldb;
+    int ldc;
+    int strideA;
+    int strideB;
+    int strideC;
+    int transA;
+    int transB;
+    int pad;
+    float2 alpha;
+    float2 beta;
+    int batch_count;
+    int pad2;
+};
+
+// -----------------------------------------------------------------------------
+// Batched Complex GEMM Kernel (c64)
+// -----------------------------------------------------------------------------
+kernel void batch_cgemm_c64(
+    device const float2* A               [[buffer(0)]],
+    device const float2* B               [[buffer(1)]],
+    device float2*       C               [[buffer(2)]],
+    constant BatchGemmParams& params     [[buffer(3)]],
+    uint3 thread_pos_in_grid             [[thread_position_in_grid]],
+    uint3 thread_pos_in_tg               [[thread_position_in_threadgroup]],
+    uint3 tg_pos                         [[threadgroup_position_in_grid]]
+) {
+    int col = thread_pos_in_grid.x;
+    int row = thread_pos_in_grid.y;
+    int batch_idx = thread_pos_in_grid.z;
+
+    if (batch_idx >= params.batch_count) {
+        return;
+    }
+
+    threadgroup float2 sA[TILE_SIZE][TILE_SIZE];
+    threadgroup float2 sB[TILE_SIZE][TILE_SIZE];
+
+    int tx = thread_pos_in_tg.x;
+    int ty = thread_pos_in_tg.y;
+
+    device const float2* currA = A + (int64_t)batch_idx * params.strideA;
+    device const float2* currB = B + (int64_t)batch_idx * params.strideB;
+    device float2*       currC = C + (int64_t)batch_idx * params.strideC;
+
+    float2 acc = float2(0.0f, 0.0f);
+
+    int num_tiles = (params.K + TILE_SIZE - 1) / TILE_SIZE;
+
+    for (int t = 0; t < num_tiles; t++) {
+        int a_k = t * TILE_SIZE + tx;
+        if (row < params.M && a_k < params.K) {
+            float2 val;
+            if (params.transA == 0) {
+                val = currA[row * params.lda + a_k];
+            } else if (params.transA == 1) {
+                val = currA[a_k * params.lda + row];
+            } else {
+                val = c_conj(currA[a_k * params.lda + row]);
+            }
+            sA[ty][tx] = val;
+        } else {
+            sA[ty][tx] = float2(0.0f, 0.0f);
+        }
+
+        int b_k = t * TILE_SIZE + ty;
+        if (b_k < params.K && col < params.N) {
+            float2 val;
+            if (params.transB == 0) {
+                val = currB[b_k * params.ldb + col];
+            } else if (params.transB == 1) {
+                val = currB[col * params.ldb + b_k];
+            } else {
+                val = c_conj(currB[col * params.ldb + b_k]);
+            }
+            sB[ty][tx] = val;
+        } else {
+            sB[ty][tx] = float2(0.0f, 0.0f);
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        #pragma unroll
+        for (int k = 0; k < TILE_SIZE; k++) {
+            acc = c_fma(sA[ty][k], sB[k][tx], acc);
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+
+    if (row < params.M && col < params.N) {
+        int c_idx = row * params.ldc + col;
+        float2 out_val = c_mul(params.alpha, acc);
+        if (params.beta.x != 0.0f || params.beta.y != 0.0f) {
+            out_val += c_mul(params.beta, currC[c_idx]);
+        }
+        currC[c_idx] = out_val;
+    }
+}
+
+// -----------------------------------------------------------------------------
+// Matrix Operations for GPU-Resident Expm & Adjoint
+// -----------------------------------------------------------------------------
+
+struct ElementwiseParams {
+    int total_elements;
+    int n;
+    int batch_count;
+    float scalar_a;
+    float scalar_b;
+    float scalar_c;
+};
+
+kernel void cmat_scale_kernel(
+    device const float2* src [[buffer(0)]],
+    device float2* dst       [[buffer(1)]],
+    constant ElementwiseParams& p [[buffer(2)]],
+    uint id [[thread_position_in_grid]]
+) {
+    if (id < (uint)p.total_elements) {
+        dst[id] = src[id] * p.scalar_a;
+    }
+}
+
+kernel void cmat_scale_adaptive_kernel(
+    device const float2* src     [[buffer(0)]],
+    device float2* dst           [[buffer(1)]],
+    device const float* scales   [[buffer(2)]],
+    constant ElementwiseParams& p [[buffer(3)]],
+    uint id [[thread_position_in_grid]]
+) {
+    if (id < (uint)p.total_elements) {
+        int b = id / (p.n * p.n);
+        dst[id] = src[id] * scales[b];
+    }
+}
+
+kernel void cmat_squaring_gate_kernel(
+    device const float2* gemm_res [[buffer(0)]],
+    device const float2* acc_prev [[buffer(1)]],
+    device float2* dst            [[buffer(2)]],
+    device const int* squarings   [[buffer(3)]],
+    constant int& current_k       [[buffer(4)]],
+    constant ElementwiseParams& p [[buffer(5)]],
+    uint id [[thread_position_in_grid]]
+) {
+    if (id < (uint)p.total_elements) {
+        int b = id / (p.n * p.n);
+        if (current_k < squarings[b]) {
+            dst[id] = gemm_res[id];
+        } else {
+            dst[id] = acc_prev[id];
+        }
+    }
+}
+
+kernel void cmat_copy_kernel(
+    device const float2* src [[buffer(0)]],
+    device float2* dst       [[buffer(1)]],
+    constant ElementwiseParams& p [[buffer(2)]],
+    uint id [[thread_position_in_grid]]
+) {
+    if (id < (uint)p.total_elements) {
+        dst[id] = src[id];
+    }
+}
+
+kernel void cmat_conjugate_transpose_kernel(
+    device const float2* src [[buffer(0)]],
+    device float2* dst       [[buffer(1)]],
+    constant ElementwiseParams& p [[buffer(2)]],
+    uint3 pos [[thread_position_in_grid]]
+) {
+    int col = pos.x;
+    int row = pos.y;
+    int b = pos.z;
+    if (col < p.n && row < p.n && b < p.batch_count) {
+        int src_idx = b * (p.n * p.n) + row * p.n + col;
+        int dst_idx = b * (p.n * p.n) + col * p.n + row;
+        float2 val = src[src_idx];
+        dst[dst_idx] = float2(val.x, -val.y);
+    }
+}
+
+kernel void cmat_taylor_horner_combine(
+    device float2* acc       [[buffer(0)]],
+    device const float2* m   [[buffer(1)]],
+    device const float2* m2  [[buffer(2)]],
+    constant ElementwiseParams& p [[buffer(3)]],
+    uint3 pos [[thread_position_in_grid]]
+) {
+    int col = pos.x;
+    int row = pos.y;
+    int b = pos.z;
+    if (col < p.n && row < p.n && b < p.batch_count) {
+        int idx = b * (p.n * p.n) + row * p.n + col;
+        float2 val = acc[idx] + m[idx] * p.scalar_b + m2[idx] * p.scalar_c;
+        if (row == col) {
+            val.x += p.scalar_a;
+        }
+        acc[idx] = val;
+    }
+}
+
+kernel void cmat_taylor_horner_combine_pair(
+    device float2* acc_y     [[buffer(0)]],
+    device float2* acc_l     [[buffer(1)]],
+    device const float2* y   [[buffer(2)]],
+    device const float2* l   [[buffer(3)]],
+    device const float2* y2  [[buffer(4)]],
+    device const float2* l2  [[buffer(5)]],
+    constant ElementwiseParams& p [[buffer(6)]],
+    uint3 pos [[thread_position_in_grid]]
+) {
+    int col = pos.x;
+    int row = pos.y;
+    int b = pos.z;
+    if (col < p.n && row < p.n && b < p.batch_count) {
+        int idx = b * (p.n * p.n) + row * p.n + col;
+        float2 vy = acc_y[idx] + y[idx] * p.scalar_b + y2[idx] * p.scalar_c;
+        if (row == col) {
+            vy.x += p.scalar_a;
+        }
+        acc_y[idx] = vy;
+        acc_l[idx] = acc_l[idx] + l[idx] * p.scalar_b + l2[idx] * p.scalar_c;
+    }
+}
+
+// -----------------------------------------------------------------------------
+// Structure Factors on Metal (c64 / float2)
+// -----------------------------------------------------------------------------
+
+struct StructureFactorParams {
+    int n_grid;
+    int n_atoms;
+    float volume;
+    float two_pi;
+    float minus_two_pi_sq;
+    int absorption;
+};
+
+static inline float lobato_eval(device const float* a_coeffs, device const float* b_coeffs, int offset, float g2) {
+    float sum = 0.0f;
+    for (int i = 0; i < 5; i++) {
+        float a = a_coeffs[offset + i];
+        float b = b_coeffs[offset + i];
+        sum += a / (1.0f + b * g2);
+    }
+    return sum;
+}
+
+kernel void structure_factors_c64_kernel(
+    device float2* fgb                        [[buffer(0)]],
+    device const float* g2_arr                [[buffer(1)]],
+    device const float* hkl_arr               [[buffer(2)]],
+    device const float* lobato_a              [[buffer(3)]],
+    device const float* lobato_b              [[buffer(4)]],
+    device const float* uij_arr               [[buffer(5)]],
+    device const float* pos_arr               [[buffer(6)]],
+    device const float* occ_arr               [[buffer(7)]],
+    constant StructureFactorParams& p         [[buffer(8)]],
+    uint id [[thread_position_in_grid]]
+) {
+    if (id >= (uint)p.n_grid) return;
+
+    float g2 = g2_arr[id];
+    float h0 = hkl_arr[id * 3 + 0];
+    float h1 = hkl_arr[id * 3 + 1];
+    float h2 = hkl_arr[id * 3 + 2];
+
+    float re_acc = 0.0f;
+    float im_acc = 0.0f;
+
+    for (int a = 0; a < p.n_atoms; a++) {
+        float form = lobato_eval(lobato_a, lobato_b, a * 5, g2);
+
+        int base_u = a * 9;
+        float quad = h0 * (uij_arr[base_u + 0] * h0 + uij_arr[base_u + 1] * h1 + uij_arr[base_u + 2] * h2)
+                   + h1 * (uij_arr[base_u + 3] * h0 + uij_arr[base_u + 4] * h1 + uij_arr[base_u + 5] * h2)
+                   + h2 * (uij_arr[base_u + 6] * h0 + uij_arr[base_u + 7] * h1 + uij_arr[base_u + 8] * h2);
+        float dwf = metal::exp(p.minus_two_pi_sq * quad);
+
+        float amplitude = (form * dwf) / p.volume;
+        float phase = p.two_pi * (pos_arr[a * 3 + 0] * h0 + pos_arr[a * 3 + 1] * h1 + pos_arr[a * 3 + 2] * h2);
+
+        float cp = metal::cos(phase);
+        float sp = metal::sin(phase);
+
+        float occ = occ_arr[a];
+        re_acc += occ * amplitude * cp;
+        im_acc += occ * amplitude * sp;
+    }
+
+    fgb[id] = float2(re_acc, im_acc);
+}
+
+// -----------------------------------------------------------------------------
+// Structure Matrix Assembly on Metal (c64 / float2)
+// -----------------------------------------------------------------------------
+
+struct StructureMatrixParams {
+    int n_beams;
+    int n_batch;
+    int buffer_size;
+    int n_grid;
+    float prefactor;
+    float u0_prime;
+};
+
+kernel void structure_matrix_scatter_kernel(
+    device float2* buffer                     [[buffer(0)]],
+    device const float2* fgb                  [[buffer(1)]],
+    device const int* source                  [[buffer(2)]],
+    constant StructureMatrixParams& p         [[buffer(3)]],
+    uint id [[thread_position_in_grid]]
+) {
+    if (id < (uint)p.n_grid) {
+        int slot = source[id];
+        buffer[slot] = fgb[id];
+    }
+}
+
+kernel void structure_matrix_gather_kernel(
+    device float2* a                          [[buffer(0)]],
+    device const float2* buffer               [[buffer(1)]],
+    device const int* destination             [[buffer(2)]],
+    device const float* mii                   [[buffer(3)]],
+    device const float* diagonal              [[buffer(4)]],
+    constant StructureMatrixParams& p         [[buffer(5)]],
+    uint3 pos [[thread_position_in_grid]]
+) {
+    int j = pos.x;
+    int i = pos.y;
+    int b = pos.z;
+
+    int n = p.n_beams;
+    if (j < n && i < n && b < p.n_batch) {
+        int base_m = b * n;
+        int base_a = b * (n * n);
+        int out_idx = base_a + i * n + j;
+
+        if (i == j) {
+            float diag_re = diagonal[base_m + i];
+            float diag_im = p.u0_prime * mii[base_m + i];
+            a[out_idx] = float2(diag_re, diag_im);
+        } else {
+            int slot = destination[i * n + j];
+            float2 z = buffer[slot];
+            float scale = p.prefactor * mii[base_m + j] * mii[base_m + i];
+            a[out_idx] = float2(z.x * scale, z.y * scale);
+        }
+    }
+}