faster-diffbloch 0.1.0__py3-none-any.whl

faster_diffbloch/native/metal_batch_cgemm.m

@@ -0,0 +1,776 @@
+#import <Foundation/Foundation.h>
+#import <Metal/Metal.h>
+#include "batch_cgemm.h"
+#include <math.h>
+#include <dlfcn.h>
+
+struct __attribute__((aligned(8))) BatchGemmParams {
+    int32_t M;
+    int32_t N;
+    int32_t K;
+    int32_t lda;
+    int32_t ldb;
+    int32_t ldc;
+    int32_t strideA;
+    int32_t strideB;
+    int32_t strideC;
+    int32_t transA;
+    int32_t transB;
+    int32_t pad;
+    struct { float x; float y; } alpha;
+    struct { float x; float y; } beta;
+    int32_t batch_count;
+    int32_t pad2;
+};
+
+struct __attribute__((aligned(8))) ElementwiseParams {
+    int32_t total_elements;
+    int32_t n;
+    int32_t batch_count;
+    float scalar_a;
+    float scalar_b;
+    float scalar_c;
+};
+
+struct __attribute__((aligned(8))) StructureFactorParams {
+    int32_t n_grid;
+    int32_t n_atoms;
+    float volume;
+    float two_pi;
+    float minus_two_pi_sq;
+    int32_t absorption;
+};
+
+struct __attribute__((aligned(8))) StructureMatrixParams {
+    int32_t n_beams;
+    int32_t n_batch;
+    int32_t buffer_size;
+    int32_t n_grid;
+    float prefactor;
+    float u0_prime;
+};
+
+static id<MTLDevice> g_device = nil;
+static id<MTLCommandQueue> g_queue = nil;
+
+static id<MTLComputePipelineState> g_pipe_gemm = nil;
+static id<MTLComputePipelineState> g_pipe_scale = nil;
+static id<MTLComputePipelineState> g_pipe_scale_adaptive = nil;
+static id<MTLComputePipelineState> g_pipe_squaring_gate = nil;
+static id<MTLComputePipelineState> g_pipe_copy = nil;
+static id<MTLComputePipelineState> g_pipe_conj_trans = nil;
+static id<MTLComputePipelineState> g_pipe_horner = nil;
+static id<MTLComputePipelineState> g_pipe_horner_pair = nil;
+static id<MTLComputePipelineState> g_pipe_sf = nil;
+static id<MTLComputePipelineState> g_pipe_sm_scatter = nil;
+static id<MTLComputePipelineState> g_pipe_sm_gather = nil;
+
+static dispatch_once_t g_init_once;
+
+static const double TAYLOR_COEFFS[19] = {
+    1.0,
+    1.0,
+    0.5,
+    1.0 / 6.0,
+    1.0 / 24.0,
+    1.0 / 120.0,
+    1.0 / 720.0,
+    1.0 / 5040.0,
+    1.0 / 40320.0,
+    1.0 / 362880.0,
+    1.0 / 3628800.0,
+    1.0 / 39916800.0,
+    1.0 / 479001600.0,
+    1.0 / 6227020800.0,
+    1.0 / 87178291200.0,
+    1.0 / 1307674368000.0,
+    1.0 / 20922789888000.0,
+    1.0 / 355687428096000.0,
+    1.0 / 6402373705728000.0
+};
+
+static const double THETA18 = 3.010066362817634;
+
+static id<MTLComputePipelineState> make_pipeline(id<MTLLibrary> lib, NSString* name, NSError** err) {
+    id<MTLFunction> fn = [lib newFunctionWithName:name];
+    if (!fn) {
+        fprintf(stderr, "Metal error: kernel function %s not found\n", name.UTF8String);
+        return nil;
+    }
+    return [g_device newComputePipelineStateWithFunction:fn error:err];
+}
+
+static NSString* get_metallib_path(void) {
+    Dl_info info;
+    if (dladdr((const void*)get_metallib_path, &info) && info.dli_fname) {
+        NSString* dylibPath = [NSString stringWithUTF8String:info.dli_fname];
+        NSString* dir = [dylibPath stringByDeletingLastPathComponent];
+        NSString* metallibInDir = [dir stringByAppendingPathComponent:@"batch_cgemm.metallib"];
+        if ([[NSFileManager defaultManager] fileExistsAtPath:metallibInDir]) {
+            return metallibInDir;
+        }
+    }
+    NSArray* candidates = @[
+        @"bench/build/batch_cgemm.metallib",
+        @"src/batch_gemm/batch_cgemm.metallib",
+        @"../bench/build/batch_cgemm.metallib",
+        @"../../bench/build/batch_cgemm.metallib",
+        @"/Users/abhishekshivakumar/diffFlow/bench/build/batch_cgemm.metallib"
+    ];
+    for (NSString* p in candidates) {
+        if ([[NSFileManager defaultManager] fileExistsAtPath:p]) {
+            return p;
+        }
+    }
+    return nil;
+}
+
+static void init_metal_state(void) {
+    dispatch_once(&g_init_once, ^{
+        @autoreleasepool {
+            g_device = MTLCreateSystemDefaultDevice();
+            if (!g_device) {
+                fprintf(stderr, "Metal error: failed to get default Metal device\n");
+                return;
+            }
+            g_queue = [g_device newCommandQueue];
+
+            NSError* error = nil;
+            NSString* libPath = get_metallib_path();
+            id<MTLLibrary> library = nil;
+            if (libPath) {
+                library = [g_device newLibraryWithURL:[NSURL fileURLWithPath:libPath] error:&error];
+            }
+
+            if (!library) {
+                fprintf(stderr, "Metal library error: %s\n", error ? error.localizedDescription.UTF8String : "library not found");
+                return;
+            }
+
+            g_pipe_gemm = make_pipeline(library, @"batch_cgemm_c64", &error);
+            g_pipe_scale = make_pipeline(library, @"cmat_scale_kernel", &error);
+            g_pipe_scale_adaptive = make_pipeline(library, @"cmat_scale_adaptive_kernel", &error);
+            g_pipe_squaring_gate = make_pipeline(library, @"cmat_squaring_gate_kernel", &error);
+            g_pipe_copy = make_pipeline(library, @"cmat_copy_kernel", &error);
+            g_pipe_conj_trans = make_pipeline(library, @"cmat_conjugate_transpose_kernel", &error);
+            g_pipe_horner = make_pipeline(library, @"cmat_taylor_horner_combine", &error);
+            g_pipe_horner_pair = make_pipeline(library, @"cmat_taylor_horner_combine_pair", &error);
+            g_pipe_sf = make_pipeline(library, @"structure_factors_c64_kernel", &error);
+            g_pipe_sm_scatter = make_pipeline(library, @"structure_matrix_scatter_kernel", &error);
+            g_pipe_sm_gather = make_pipeline(library, @"structure_matrix_gather_kernel", &error);
+        }
+    });
+}
+
+static inline int metal_trans_mode(BatchGemmTranspose t) {
+    switch (t) {
+        case BATCH_CGEMM_NO_TRANS:   return 0;
+        case BATCH_CGEMM_TRANS:      return 1;
+        case BATCH_CGEMM_CONJ_TRANS: return 2;
+        default:                     return 0;
+    }
+}
+
+static void encode_cgemm(
+    id<MTLComputeCommandEncoder> enc,
+    id<MTLBuffer> bufA, int64_t offA, int strideA, int transA,
+    id<MTLBuffer> bufB, int64_t offB, int strideB, int transB,
+    id<MTLBuffer> bufC, int64_t offC, int strideC,
+    int M, int N, int K,
+    float alpha_r, float alpha_i,
+    float beta_r, float beta_i,
+    int batch_count
+) {
+    struct BatchGemmParams params;
+    params.M = M;
+    params.N = N;
+    params.K = K;
+    params.lda = (transA == 0) ? K : M;
+    params.ldb = (transB == 0) ? N : K;
+    params.ldc = N;
+    params.strideA = strideA;
+    params.strideB = strideB;
+    params.strideC = strideC;
+    params.transA = transA;
+    params.transB = transB;
+    params.pad = 0;
+    params.alpha.x = alpha_r;
+    params.alpha.y = alpha_i;
+    params.beta.x = beta_r;
+    params.beta.y = beta_i;
+    params.batch_count = batch_count;
+    params.pad2 = 0;
+
+    [enc setComputePipelineState:g_pipe_gemm];
+    [enc setBuffer:bufA offset:offA atIndex:0];
+    [enc setBuffer:bufB offset:offB atIndex:1];
+    [enc setBuffer:bufC offset:offC atIndex:2];
+    [enc setBytes:&params length:sizeof(params) atIndex:3];
+
+    MTLSize threadsPerTG = MTLSizeMake(16, 16, 1);
+    MTLSize tgPerGrid = MTLSizeMake((N + 15) / 16, (M + 15) / 16, batch_count);
+    [enc dispatchThreadgroups:tgPerGrid threadsPerThreadgroup:threadsPerTG];
+    [enc memoryBarrierWithScope:MTLBarrierScopeBuffers];
+}
+
+static double matrix_norm1_c64(const complex64_t* A, int n) {
+    double worst = 0.0;
+    for (int j = 0; j < n; j++) {
+        double col_sum = 0.0;
+        for (int i = 0; i < n; i++) {
+            col_sum += cabsf(A[i * n + j]);
+        }
+        if (col_sum > worst) worst = col_sum;
+    }
+    return worst;
+}
+
+static double pair_norm1_c64(const complex64_t* Y, const complex64_t* L, int n) {
+    double worst = 0.0;
+    for (int j = 0; j < n; j++) {
+        double col_y = 0.0;
+        double col_l = 0.0;
+        for (int i = 0; i < n; i++) {
+            col_y += cabsf(Y[i * n + j]);
+            col_l += cabsf(L[i * n + j]);
+        }
+        if (col_y > worst) worst = col_y;
+        if (col_y + col_l > worst) worst = col_y + col_l;
+    }
+    return worst;
+}
+
+// -----------------------------------------------------------------------------
+// Strided Batched Complex GEMM (c64)
+// -----------------------------------------------------------------------------
+int32_t batch_cgemm_strided_metal_c64(
+    BatchGemmTranspose transA,
+    BatchGemmTranspose transB,
+    int32_t M,
+    int32_t N,
+    int32_t K,
+    const complex64_t* alpha,
+    const complex64_t* A,
+    int32_t lda,
+    int64_t strideA,
+    const complex64_t* B,
+    int32_t ldb,
+    int64_t strideB,
+    const complex64_t* beta,
+    complex64_t* C,
+    int32_t ldc,
+    int64_t strideC,
+    int32_t batch_count
+) {
+    if (batch_count <= 0) return 0;
+    init_metal_state();
+    if (!g_pipe_gemm || !g_queue || !g_device) return -1;
+
+    @autoreleasepool {
+        size_t bytesA = (size_t)batch_count * strideA * sizeof(complex64_t);
+        size_t bytesB = (size_t)batch_count * strideB * sizeof(complex64_t);
+        size_t bytesC = (size_t)batch_count * strideC * sizeof(complex64_t);
+
+        id<MTLBuffer> bufA = [g_device newBufferWithBytes:A length:bytesA options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufB = [g_device newBufferWithBytes:B length:bytesB options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufC = [g_device newBufferWithBytes:C length:bytesC options:MTLResourceStorageModeShared];
+
+        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
+        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
+
+        encode_cgemm(
+            enc,
+            bufA, 0, (int)strideA, metal_trans_mode(transA),
+            bufB, 0, (int)strideB, metal_trans_mode(transB),
+            bufC, 0, (int)strideC,
+            M, N, K,
+            crealf(*alpha), cimagf(*alpha),
+            crealf(*beta), cimagf(*beta),
+            batch_count
+        );
+
+        [enc endEncoding];
+        [cmd commit];
+        [cmd waitUntilCompleted];
+
+        memcpy(C, bufC.contents, bytesC);
+    }
+    return 0;
+}
+
+// -----------------------------------------------------------------------------
+// Metal Batched Matrix Exponential (c64)
+// -----------------------------------------------------------------------------
+int32_t metal_matrix_exp_c64(
+    const complex64_t* A,
+    int32_t n,
+    int32_t batch_count,
+    complex64_t* out
+) {
+    if (batch_count <= 0 || n <= 0) return 0;
+    init_metal_state();
+    if (!g_pipe_gemm || !g_queue || !g_device) return -1;
+
+    @autoreleasepool {
+        int matrix_elements = n * n;
+        int total_elements = batch_count * matrix_elements;
+        size_t buf_bytes = (size_t)total_elements * sizeof(complex64_t);
+
+        int max_squarings = 0;
+        int* squarings = (int*)malloc((size_t)batch_count * sizeof(int));
+        float* scales = (float*)malloc((size_t)batch_count * sizeof(float));
+
+        for (int b = 0; b < batch_count; b++) {
+            double norm = matrix_norm1_c64(A + b * matrix_elements, n);
+            int sq = 0;
+            if (norm > THETA18) {
+                sq = (int)ceil(log2(norm / THETA18));
+            }
+            squarings[b] = sq;
+            scales[b] = (sq > 0) ? (float)exp(-sq * log(2.0)) : 1.0f;
+            if (sq > max_squarings) max_squarings = sq;
+        }
+
+        id<MTLBuffer> bufM         = [g_device newBufferWithBytes:A length:buf_bytes options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufM2        = [g_device newBufferWithLength:buf_bytes options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufM3        = [g_device newBufferWithLength:buf_bytes options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufAcc       = [g_device newBufferWithLength:buf_bytes options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufTmp       = [g_device newBufferWithLength:buf_bytes options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufScales    = [g_device newBufferWithBytes:scales length:(size_t)batch_count * sizeof(float) options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufSquarings = [g_device newBufferWithBytes:squarings length:(size_t)batch_count * sizeof(int) options:MTLResourceStorageModeShared];
+
+        free(squarings);
+        free(scales);
+
+        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
+        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
+
+        struct ElementwiseParams ep;
+        ep.total_elements = total_elements;
+        ep.n = n;
+        ep.batch_count = batch_count;
+        ep.scalar_a = 0.0f;
+        ep.scalar_b = 0.0f;
+        ep.scalar_c = 0.0f;
+
+        // 1. Adaptive per-matrix scaling
+        [enc setComputePipelineState:g_pipe_scale_adaptive];
+        [enc setBuffer:bufM offset:0 atIndex:0];
+        [enc setBuffer:bufM offset:0 atIndex:1];
+        [enc setBuffer:bufScales offset:0 atIndex:2];
+        [enc setBytes:&ep length:sizeof(ep) atIndex:3];
+        [enc dispatchThreads:MTLSizeMake(total_elements, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
+        [enc memoryBarrierWithScope:MTLBarrierScopeBuffers];
+
+        // 2. M2 = M * M, M3 = M2 * M
+        encode_cgemm(enc, bufM, 0, matrix_elements, 0, bufM, 0, matrix_elements, 0, bufM2, 0, matrix_elements, n, n, n, 1.0f, 0.0f, 0.0f, 0.0f, batch_count);
+        encode_cgemm(enc, bufM2, 0, matrix_elements, 0, bufM, 0, matrix_elements, 0, bufM3, 0, matrix_elements, n, n, n, 1.0f, 0.0f, 0.0f, 0.0f, batch_count);
+
+        // 3. Horner evaluation in M3
+        for (int j = 5; j >= 0; j--) {
+            if (j < 5) {
+                encode_cgemm(enc, bufAcc, 0, matrix_elements, 0, bufM3, 0, matrix_elements, 0, bufTmp, 0, matrix_elements, n, n, n, 1.0f, 0.0f, 0.0f, 0.0f, batch_count);
+                [enc setComputePipelineState:g_pipe_copy];
+                [enc setBuffer:bufTmp offset:0 atIndex:0];
+                [enc setBuffer:bufAcc offset:0 atIndex:1];
+                [enc setBytes:&ep length:sizeof(ep) atIndex:2];
+                [enc dispatchThreads:MTLSizeMake(total_elements, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
+                [enc memoryBarrierWithScope:MTLBarrierScopeBuffers];
+            } else {
+                ep.scalar_a = 0.0f;
+                [enc setComputePipelineState:g_pipe_scale];
+                [enc setBuffer:bufAcc offset:0 atIndex:0];
+                [enc setBuffer:bufAcc offset:0 atIndex:1];
+                [enc setBytes:&ep length:sizeof(ep) atIndex:2];
+                [enc dispatchThreads:MTLSizeMake(total_elements, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
+                [enc memoryBarrierWithScope:MTLBarrierScopeBuffers];
+            }
+
+            ep.scalar_a = (float)TAYLOR_COEFFS[j * 3 + 0];
+            ep.scalar_b = (float)TAYLOR_COEFFS[j * 3 + 1];
+            ep.scalar_c = (float)TAYLOR_COEFFS[j * 3 + 2];
+
+            [enc setComputePipelineState:g_pipe_horner];
+            [enc setBuffer:bufAcc offset:0 atIndex:0];
+            [enc setBuffer:bufM offset:0 atIndex:1];
+            [enc setBuffer:bufM2 offset:0 atIndex:2];
+            [enc setBytes:&ep length:sizeof(ep) atIndex:3];
+            MTLSize grid = MTLSizeMake(n, n, batch_count);
+            MTLSize tg = MTLSizeMake(16, 16, 1);
+            [enc dispatchThreads:grid threadsPerThreadgroup:tg];
+            [enc memoryBarrierWithScope:MTLBarrierScopeBuffers];
+        }
+
+        // 4. Squaring loop with adaptive gating
+        for (int k = 0; k < max_squarings; k++) {
+            encode_cgemm(enc, bufAcc, 0, matrix_elements, 0, bufAcc, 0, matrix_elements, 0, bufTmp, 0, matrix_elements, n, n, n, 1.0f, 0.0f, 0.0f, 0.0f, batch_count);
+            int current_k = k;
+            [enc setComputePipelineState:g_pipe_squaring_gate];
+            [enc setBuffer:bufTmp offset:0 atIndex:0];
+            [enc setBuffer:bufAcc offset:0 atIndex:1];
+            [enc setBuffer:bufAcc offset:0 atIndex:2];
+            [enc setBuffer:bufSquarings offset:0 atIndex:3];
+            [enc setBytes:&current_k length:sizeof(int) atIndex:4];
+            [enc setBytes:&ep length:sizeof(ep) atIndex:5];
+            [enc dispatchThreads:MTLSizeMake(total_elements, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
+            [enc memoryBarrierWithScope:MTLBarrierScopeBuffers];
+        }
+
+        [enc endEncoding];
+        [cmd commit];
+        [cmd waitUntilCompleted];
+
+        memcpy(out, bufAcc.contents, buf_bytes);
+    }
+    return 0;
+}
+
+// -----------------------------------------------------------------------------
+// Metal Batched Matrix Exponential Adjoint (c64)
+// -----------------------------------------------------------------------------
+int32_t metal_matrix_exp_backward_c64(
+    const complex64_t* M,
+    const complex64_t* Ebar,
+    int32_t n,
+    int32_t batch_count,
+    int32_t dense,
+    complex64_t* Mbar
+) {
+    if (batch_count <= 0 || n <= 0) return 0;
+    init_metal_state();
+    if (!g_pipe_gemm || !g_queue || !g_device) return -1;
+
+    @autoreleasepool {
+        int matrix_elements = n * n;
+        int total_elements = batch_count * matrix_elements;
+        size_t buf_bytes = (size_t)total_elements * sizeof(complex64_t);
+
+        complex64_t* host_Y = (complex64_t*)malloc(buf_bytes);
+        for (int b = 0; b < batch_count; b++) {
+            const complex64_t* cur_m = M + b * matrix_elements;
+            complex64_t* cur_y = host_Y + b * matrix_elements;
+            for (int i = 0; i < n; i++) {
+                for (int j = 0; j < n; j++) {
+                    complex64_t z = cur_m[j * n + i];
+                    cur_y[i * n + j] = crealf(z) - I * cimagf(z);
+                }
+            }
+        }
+
+        int max_squarings = 0;
+        int* squarings = (int*)malloc((size_t)batch_count * sizeof(int));
+        float* scales = (float*)malloc((size_t)batch_count * sizeof(float));
+
+        for (int b = 0; b < batch_count; b++) {
+            double pnorm = pair_norm1_c64(host_Y + b * matrix_elements, Ebar + b * matrix_elements, n);
+            int sq = 0;
+            if (pnorm > THETA18) {
+                sq = (int)ceil(log2(pnorm / THETA18));
+            }
+            squarings[b] = sq;
+            scales[b] = (sq > 0) ? (float)exp(-sq * log(2.0)) : 1.0f;
+            if (sq > max_squarings) max_squarings = sq;
+        }
+
+        id<MTLBuffer> bufY         = [g_device newBufferWithBytes:host_Y length:buf_bytes options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufL         = [g_device newBufferWithBytes:Ebar length:buf_bytes options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufY2        = [g_device newBufferWithLength:buf_bytes options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufL2        = [g_device newBufferWithLength:buf_bytes options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufY3        = [g_device newBufferWithLength:buf_bytes options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufL3        = [g_device newBufferWithLength:buf_bytes options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufAccY      = [g_device newBufferWithLength:buf_bytes options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufAccL      = [g_device newBufferWithLength:buf_bytes options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufTmpY      = [g_device newBufferWithLength:buf_bytes options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufTmpL      = [g_device newBufferWithLength:buf_bytes options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufScales    = [g_device newBufferWithBytes:scales length:(size_t)batch_count * sizeof(float) options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufSquarings = [g_device newBufferWithBytes:squarings length:(size_t)batch_count * sizeof(int) options:MTLResourceStorageModeShared];
+
+        free(host_Y);
+        free(squarings);
+        free(scales);
+
+        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
+        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
+
+        struct ElementwiseParams ep;
+        ep.total_elements = total_elements;
+        ep.n = n;
+        ep.batch_count = batch_count;
+        ep.scalar_a = 0.0f;
+        ep.scalar_b = 0.0f;
+        ep.scalar_c = 0.0f;
+
+        // 1. Adaptive scaling of Y and L
+        [enc setComputePipelineState:g_pipe_scale_adaptive];
+        [enc setBuffer:bufY offset:0 atIndex:0];
+        [enc setBuffer:bufY offset:0 atIndex:1];
+        [enc setBuffer:bufScales offset:0 atIndex:2];
+        [enc setBytes:&ep length:sizeof(ep) atIndex:3];
+        [enc dispatchThreads:MTLSizeMake(total_elements, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
+
+        [enc setBuffer:bufL offset:0 atIndex:0];
+        [enc setBuffer:bufL offset:0 atIndex:1];
+        [enc setBuffer:bufScales offset:0 atIndex:2];
+        [enc setBytes:&ep length:sizeof(ep) atIndex:3];
+        [enc dispatchThreads:MTLSizeMake(total_elements, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
+        [enc memoryBarrierWithScope:MTLBarrierScopeBuffers];
+
+        void (^encode_pair_mul)(id<MTLBuffer>, id<MTLBuffer>, id<MTLBuffer>, id<MTLBuffer>, id<MTLBuffer>, id<MTLBuffer>) =
+            ^(id<MTLBuffer> y1, id<MTLBuffer> l1, id<MTLBuffer> y2, id<MTLBuffer> l2, id<MTLBuffer> yo, id<MTLBuffer> lo) {
+                encode_cgemm(enc, y1, 0, matrix_elements, 0, y2, 0, matrix_elements, 0, yo, 0, matrix_elements, n, n, n, 1.0f, 0.0f, 0.0f, 0.0f, batch_count);
+                encode_cgemm(enc, y1, 0, matrix_elements, 0, l2, 0, matrix_elements, 0, lo, 0, matrix_elements, n, n, n, 1.0f, 0.0f, 0.0f, 0.0f, batch_count);
+                encode_cgemm(enc, l1, 0, matrix_elements, 0, y2, 0, matrix_elements, 0, lo, 0, matrix_elements, n, n, n, 1.0f, 0.0f, 1.0f, 0.0f, batch_count);
+            };
+
+        encode_pair_mul(bufY, bufL, bufY, bufL, bufY2, bufL2);
+        encode_pair_mul(bufY2, bufL2, bufY, bufL, bufY3, bufL3);
+
+        for (int j = 5; j >= 0; j--) {
+            if (j < 5) {
+                encode_pair_mul(bufAccY, bufAccL, bufY3, bufL3, bufTmpY, bufTmpL);
+                [enc setComputePipelineState:g_pipe_copy];
+                [enc setBuffer:bufTmpY offset:0 atIndex:0];
+                [enc setBuffer:bufAccY offset:0 atIndex:1];
+                [enc setBytes:&ep length:sizeof(ep) atIndex:2];
+                [enc dispatchThreads:MTLSizeMake(total_elements, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
+
+                [enc setBuffer:bufTmpL offset:0 atIndex:0];
+                [enc setBuffer:bufAccL offset:0 atIndex:1];
+                [enc setBytes:&ep length:sizeof(ep) atIndex:2];
+                [enc dispatchThreads:MTLSizeMake(total_elements, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
+                [enc memoryBarrierWithScope:MTLBarrierScopeBuffers];
+            } else {
+                ep.scalar_a = 0.0f;
+                [enc setComputePipelineState:g_pipe_scale];
+                [enc setBuffer:bufAccY offset:0 atIndex:0];
+                [enc setBuffer:bufAccY offset:0 atIndex:1];
+                [enc setBytes:&ep length:sizeof(ep) atIndex:2];
+                [enc dispatchThreads:MTLSizeMake(total_elements, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
+
+                [enc setBuffer:bufAccL offset:0 atIndex:0];
+                [enc setBuffer:bufAccL offset:0 atIndex:1];
+                [enc setBytes:&ep length:sizeof(ep) atIndex:2];
+                [enc dispatchThreads:MTLSizeMake(total_elements, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
+                [enc memoryBarrierWithScope:MTLBarrierScopeBuffers];
+            }
+
+            ep.scalar_a = (float)TAYLOR_COEFFS[j * 3 + 0];
+            ep.scalar_b = (float)TAYLOR_COEFFS[j * 3 + 1];
+            ep.scalar_c = (float)TAYLOR_COEFFS[j * 3 + 2];
+
+            [enc setComputePipelineState:g_pipe_horner_pair];
+            [enc setBuffer:bufAccY offset:0 atIndex:0];
+            [enc setBuffer:bufAccL offset:0 atIndex:1];
+            [enc setBuffer:bufY offset:0 atIndex:2];
+            [enc setBuffer:bufL offset:0 atIndex:3];
+            [enc setBuffer:bufY2 offset:0 atIndex:4];
+            [enc setBuffer:bufL2 offset:0 atIndex:5];
+            [enc setBytes:&ep length:sizeof(ep) atIndex:6];
+            MTLSize grid = MTLSizeMake(n, n, batch_count);
+            MTLSize tg = MTLSizeMake(16, 16, 1);
+            [enc dispatchThreads:grid threadsPerThreadgroup:tg];
+            [enc memoryBarrierWithScope:MTLBarrierScopeBuffers];
+        }
+
+        for (int k = 0; k < max_squarings; k++) {
+            encode_pair_mul(bufAccY, bufAccL, bufAccY, bufAccL, bufTmpY, bufTmpL);
+            int current_k = k;
+            [enc setComputePipelineState:g_pipe_squaring_gate];
+            [enc setBuffer:bufTmpY offset:0 atIndex:0];
+            [enc setBuffer:bufAccY offset:0 atIndex:1];
+            [enc setBuffer:bufAccY offset:0 atIndex:2];
+            [enc setBuffer:bufSquarings offset:0 atIndex:3];
+            [enc setBytes:&current_k length:sizeof(int) atIndex:4];
+            [enc setBytes:&ep length:sizeof(ep) atIndex:5];
+            [enc dispatchThreads:MTLSizeMake(total_elements, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
+
+            [enc setBuffer:bufTmpL offset:0 atIndex:0];
+            [enc setBuffer:bufAccL offset:0 atIndex:1];
+            [enc setBuffer:bufAccL offset:0 atIndex:2];
+            [enc setBuffer:bufSquarings offset:0 atIndex:3];
+            [enc setBytes:&current_k length:sizeof(int) atIndex:4];
+            [enc setBytes:&ep length:sizeof(ep) atIndex:5];
+            [enc dispatchThreads:MTLSizeMake(total_elements, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
+            [enc memoryBarrierWithScope:MTLBarrierScopeBuffers];
+        }
+
+        [enc endEncoding];
+        [cmd commit];
+        [cmd waitUntilCompleted];
+
+        memcpy(Mbar, bufAccL.contents, buf_bytes);
+    }
+    return 0;
+}
+
+// -----------------------------------------------------------------------------
+// Metal Structure Factors (c128 in host, float compute on GPU)
+// -----------------------------------------------------------------------------
+int32_t metal_structure_factors_c128(
+    complex128_t* fgb,
+    const double* g2,
+    const double* hkl,
+    const double* lobato_a,
+    const double* lobato_b,
+    const double* uij,
+    const double* pos,
+    const double* occ,
+    int32_t n_grid,
+    int32_t n_atoms,
+    double volume
+) {
+    if (n_grid <= 0 || n_atoms <= 0) return 0;
+    init_metal_state();
+    if (!g_pipe_sf || !g_queue || !g_device) return -1;
+
+    @autoreleasepool {
+        float* f_g2 = (float*)malloc((size_t)n_grid * sizeof(float));
+        float* f_hkl = (float*)malloc((size_t)n_grid * 3 * sizeof(float));
+        for (int i = 0; i < n_grid; i++) f_g2[i] = (float)g2[i];
+        for (int i = 0; i < n_grid * 3; i++) f_hkl[i] = (float)hkl[i];
+
+        float* f_la = (float*)malloc((size_t)n_atoms * 5 * sizeof(float));
+        float* f_lb = (float*)malloc((size_t)n_atoms * 5 * sizeof(float));
+        float* f_uij = (float*)malloc((size_t)n_atoms * 9 * sizeof(float));
+        float* f_pos = (float*)malloc((size_t)n_atoms * 3 * sizeof(float));
+        float* f_occ = (float*)malloc((size_t)n_atoms * sizeof(float));
+
+        for (int i = 0; i < n_atoms * 5; i++) f_la[i] = (float)lobato_a[i];
+        for (int i = 0; i < n_atoms * 5; i++) f_lb[i] = (float)lobato_b[i];
+        for (int i = 0; i < n_atoms * 9; i++) f_uij[i] = (float)uij[i];
+        for (int i = 0; i < n_atoms * 3; i++) f_pos[i] = (float)pos[i];
+        for (int i = 0; i < n_atoms; i++) f_occ[i] = (float)occ[i];
+
+        id<MTLBuffer> bufFgb = [g_device newBufferWithLength:(size_t)n_grid * sizeof(complex64_t) options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufG2  = [g_device newBufferWithBytes:f_g2 length:(size_t)n_grid * sizeof(float) options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufHkl = [g_device newBufferWithBytes:f_hkl length:(size_t)n_grid * 3 * sizeof(float) options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufLa  = [g_device newBufferWithBytes:f_la length:(size_t)n_atoms * 5 * sizeof(float) options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufLb  = [g_device newBufferWithBytes:f_lb length:(size_t)n_atoms * 5 * sizeof(float) options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufUij = [g_device newBufferWithBytes:f_uij length:(size_t)n_atoms * 9 * sizeof(float) options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufPos = [g_device newBufferWithBytes:f_pos length:(size_t)n_atoms * 3 * sizeof(float) options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufOcc = [g_device newBufferWithBytes:f_occ length:(size_t)n_atoms * sizeof(float) options:MTLResourceStorageModeShared];
+
+        free(f_g2); free(f_hkl); free(f_la); free(f_lb); free(f_uij); free(f_pos); free(f_occ);
+
+        struct StructureFactorParams p;
+        p.n_grid = n_grid;
+        p.n_atoms = n_atoms;
+        p.volume = (float)volume;
+        p.two_pi = 6.28318530717958647692f;
+        p.minus_two_pi_sq = -19.73920880217871723766f;
+        p.absorption = 0;
+
+        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
+        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
+
+        [enc setComputePipelineState:g_pipe_sf];
+        [enc setBuffer:bufFgb offset:0 atIndex:0];
+        [enc setBuffer:bufG2 offset:0 atIndex:1];
+        [enc setBuffer:bufHkl offset:0 atIndex:2];
+        [enc setBuffer:bufLa offset:0 atIndex:3];
+        [enc setBuffer:bufLb offset:0 atIndex:4];
+        [enc setBuffer:bufUij offset:0 atIndex:5];
+        [enc setBuffer:bufPos offset:0 atIndex:6];
+        [enc setBuffer:bufOcc offset:0 atIndex:7];
+        [enc setBytes:&p length:sizeof(p) atIndex:8];
+
+        [enc dispatchThreads:MTLSizeMake(n_grid, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
+        [enc endEncoding];
+
+        [cmd commit];
+        [cmd waitUntilCompleted];
+
+        complex64_t* res_c64 = (complex64_t*)bufFgb.contents;
+        for (int i = 0; i < n_grid; i++) {
+            fgb[i] = (double)crealf(res_c64[i]) + I * (double)cimagf(res_c64[i]);
+        }
+    }
+    return 0;
+}
+
+// -----------------------------------------------------------------------------
+// Metal Structure Matrix Assembly (c128 in host, float compute on GPU)
+// -----------------------------------------------------------------------------
+int32_t metal_structure_matrix_c128(
+    const complex128_t* fgb,
+    int32_t n_grid,
+    const int32_t* source,
+    int32_t buffer_size,
+    const int32_t* destination,
+    int32_t n_beams,
+    const double* mii,
+    const double* diagonal,
+    int32_t n_batch,
+    double prefactor,
+    int32_t absorption,
+    complex128_t* out
+) {
+    if (n_beams <= 0 || n_batch <= 0 || n_grid <= 0 || buffer_size <= 0) return 0;
+    init_metal_state();
+    if (!g_pipe_sm_scatter || !g_pipe_sm_gather || !g_queue || !g_device) return -1;
+
+    @autoreleasepool {
+        int pairs = n_beams * n_beams;
+        int total_pairs = n_batch * pairs;
+
+        complex64_t* f_fgb = (complex64_t*)malloc((size_t)n_grid * sizeof(complex64_t));
+        for (int i = 0; i < n_grid; i++) f_fgb[i] = (float)creal(fgb[i]) + I * (float)cimag(fgb[i]);
+
+        float* f_mii = (float*)malloc((size_t)n_batch * n_beams * sizeof(float));
+        float* f_diag = (float*)malloc((size_t)n_batch * n_beams * sizeof(float));
+        for (int i = 0; i < n_batch * n_beams; i++) f_mii[i] = (float)mii[i];
+        for (int i = 0; i < n_batch * n_beams; i++) f_diag[i] = (float)diagonal[i];
+
+        id<MTLBuffer> bufBuffer = [g_device newBufferWithLength:(size_t)buffer_size * sizeof(complex64_t) options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufFgb    = [g_device newBufferWithBytes:f_fgb length:(size_t)n_grid * sizeof(complex64_t) options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufSource = [g_device newBufferWithBytes:source length:(size_t)n_grid * sizeof(int32_t) options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufDst    = [g_device newBufferWithBytes:destination length:(size_t)pairs * sizeof(int32_t) options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufMii    = [g_device newBufferWithBytes:f_mii length:(size_t)n_batch * n_beams * sizeof(float) options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufDiag   = [g_device newBufferWithBytes:f_diag length:(size_t)n_batch * n_beams * sizeof(float) options:MTLResourceStorageModeShared];
+        id<MTLBuffer> bufOut    = [g_device newBufferWithLength:(size_t)total_pairs * sizeof(complex64_t) options:MTLResourceStorageModeShared];
+
+        free(f_fgb); free(f_mii); free(f_diag);
+
+        struct StructureMatrixParams p;
+        p.n_beams = n_beams;
+        p.n_batch = n_batch;
+        p.buffer_size = buffer_size;
+        p.n_grid = n_grid;
+        p.prefactor = (float)prefactor;
+        p.u0_prime = 0.0f;
+
+        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
+        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
+
+        // 1. Scatter into buffer
+        [enc setComputePipelineState:g_pipe_sm_scatter];
+        [enc setBuffer:bufBuffer offset:0 atIndex:0];
+        [enc setBuffer:bufFgb offset:0 atIndex:1];
+        [enc setBuffer:bufSource offset:0 atIndex:2];
+        [enc setBytes:&p length:sizeof(p) atIndex:3];
+        [enc dispatchThreads:MTLSizeMake(n_grid, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
+        [enc memoryBarrierWithScope:MTLBarrierScopeBuffers];
+
+        // 2. Gather into A
+        [enc setComputePipelineState:g_pipe_sm_gather];
+        [enc setBuffer:bufOut offset:0 atIndex:0];
+        [enc setBuffer:bufBuffer offset:0 atIndex:1];
+        [enc setBuffer:bufDst offset:0 atIndex:2];
+        [enc setBuffer:bufMii offset:0 atIndex:3];
+        [enc setBuffer:bufDiag offset:0 atIndex:4];
+        [enc setBytes:&p length:sizeof(p) atIndex:5];
+        MTLSize grid = MTLSizeMake(n_beams, n_beams, n_batch);
+        MTLSize tg = MTLSizeMake(16, 16, 1);
+        [enc dispatchThreads:grid threadsPerThreadgroup:tg];
+
+        [enc endEncoding];
+        [cmd commit];
+        [cmd waitUntilCompleted];
+
+        complex64_t* res_c64 = (complex64_t*)bufOut.contents;
+        for (int i = 0; i < total_pairs; i++) {
+            out[i] = (double)crealf(res_c64[i]) + I * (double)cimagf(res_c64[i]);
+        }
+    }
+    return 0;
+}