// metalcore PyTorch bindings
// Provides Python interface to Metal kernels for QR, trsm, Householder
//
// Based on metalsvd pattern: 
// - Load kernels from .metal file
// - Cache pipeline states
// - Use MPS stream for synchronization

#include <torch/extension.h>
#include <ATen/ATen.h>
#include <ATen/mps/MPSDevice.h>
#include <ATen/mps/MPSStream.h>
#include <ATen/native/mps/OperationUtils.h>

#import <Metal/Metal.h>
#import <Foundation/Foundation.h>

#include <fstream>
#include <sstream>
#include <dlfcn.h>
#include <libgen.h>

using namespace at::mps;
using namespace at::native::mps;

// -----------------------------------------------------------------------------
// Global State
// -----------------------------------------------------------------------------

struct CoreKernels {
    // Panel QR (geqr2)
    id<MTLFunction> geqr2 = nil;
    id<MTLComputePipelineState> geqr2PSO = nil;
    
    // Fused Panel QR (MAGMA-style)
    id<MTLFunction> geqr2Fused = nil;
    id<MTLComputePipelineState> geqr2FusedPSO = nil;
    
    // Householder
    id<MTLFunction> householder = nil;
    id<MTLFunction> applyHouseholder = nil;
    id<MTLComputePipelineState> householderPSO = nil;
    id<MTLComputePipelineState> applyHouseholderPSO = nil;
    
    // Block reflector (larfb)
    id<MTLFunction> larfbStep1 = nil;
    id<MTLFunction> larfbStep2 = nil;
    id<MTLFunction> larfbStep3 = nil;
    id<MTLComputePipelineState> larfbStep1PSO = nil;
    id<MTLComputePipelineState> larfbStep2PSO = nil;
    id<MTLComputePipelineState> larfbStep3PSO = nil;
    
    // T matrix builder
    id<MTLFunction> larft = nil;
    id<MTLComputePipelineState> larftPSO = nil;
    
    // Triangular solve
    id<MTLFunction> trsmLower = nil;
    id<MTLFunction> trsmUpper = nil;
    id<MTLComputePipelineState> trsmLowerPSO = nil;
    id<MTLComputePipelineState> trsmUpperPSO = nil;
    
    // Fully fused QR (single dispatch)
    id<MTLFunction> qrFullFused = nil;
    id<MTLComputePipelineState> qrFullFusedPSO = nil;
    
    // Batched QR (parallel matrices)
    id<MTLFunction> qrBatched = nil;
    id<MTLComputePipelineState> qrBatchedPSO = nil;
    
    // Batched TRSM (triangular solve)
    id<MTLFunction> trsmBatched = nil;
    id<MTLComputePipelineState> trsmBatchedPSO = nil;
    
    // Column norms
    id<MTLFunction> columnNorms = nil;
    id<MTLComputePipelineState> columnNormsPSO = nil;
    
    // Batched Cholesky (potrf)
    id<MTLFunction> choleskyBatched = nil;
    id<MTLComputePipelineState> choleskyBatchedPSO = nil;
    
    // Batched TRSM lower/upper
    id<MTLFunction> trsmLowerBatched = nil;
    id<MTLFunction> trsmUpperBatched = nil;
    id<MTLComputePipelineState> trsmLowerBatchedPSO = nil;
    id<MTLComputePipelineState> trsmUpperBatchedPSO = nil;
    
    // Fused Cholesky solve (forward + back substitution in one kernel)
    id<MTLFunction> choleskySolveBatched = nil;
    id<MTLComputePipelineState> choleskySolveBatchedPSO = nil;
    
    // Column norm sort (De Rijk optimization for SVD)
    id<MTLFunction> columnNormSort = nil;
    id<MTLComputePipelineState> columnNormSortPSO = nil;
    
    // Sign canonicalization (SVD U/V sign fix)
    id<MTLFunction> signCanonicalize = nil;
    id<MTLComputePipelineState> signCanonicalizePSO = nil;
    
    // Batched Q.T @ b for fused solve
    id<MTLFunction> batchedQtB = nil;
    id<MTLComputePipelineState> batchedQtBPSO = nil;
    
    // High-impact ML/LA kernels
    id<MTLFunction> luBatched = nil;
    id<MTLComputePipelineState> luBatchedPSO = nil;
    
    id<MTLFunction> inverseBatched = nil;
    id<MTLComputePipelineState> inverseBatchedPSO = nil;
    
    id<MTLFunction> syrkBatched = nil;
    id<MTLComputePipelineState> syrkBatchedPSO = nil;
    
    id<MTLFunction> frobeniusNormBatched = nil;
    id<MTLComputePipelineState> frobeniusNormBatchedPSO = nil;
    
    id<MTLFunction> softmaxBatched = nil;
    id<MTLComputePipelineState> softmaxBatchedPSO = nil;
    
    id<MTLFunction> traceBatched = nil;
    id<MTLComputePipelineState> traceBatchedPSO = nil;
};

static CoreKernels kernels;
static id<MTLLibrary> coreLib = nil;
static std::once_flag init_flag;

// -----------------------------------------------------------------------------
// Kernel Loading
// -----------------------------------------------------------------------------

void load_core_kernels() {
    std::call_once(init_flag, [](){
        id<MTLDevice> device = MPSDevice::getInstance()->device();
        if (!device) TORCH_CHECK(false, "MPS Device not found");
        
        NSError* error = nil;
        
        // 1. Try to locate .metallib relative to this dylib
        // Get path to this dylib
        Dl_info info;
        if (dladdr((void*)&load_core_kernels, &info)) {
            // info.dli_fname contains path to metalcore_backend.so
            // Structure in wheel:
            //   site-packages/metalcore_backend.cpython...so
            //   site-packages/metalcore/native/core_kernels.metallib
            // OR if built in-place:
            //   src/metalcore_backend.cpython...so
            //   src/metalcore/native/
            
            // We need to look for metalcore/native/core_kernels.metallib relative to the directory containing the .so
            
            std::string dylib_path = info.dli_fname;
            std::string dylib_dir = dylib_path.substr(0, dylib_path.find_last_of('/'));
            
            // Try explicit path options
            std::vector<std::string> candidates = {
                dylib_dir + "/metalcore/native/core_kernels.metallib",     // Wheel structure (if backend is top level)
                dylib_dir + "/native/core_kernels.metallib",               // In-place or nested
                dylib_dir + "/../metalcore/native/core_kernels.metallib"   // Sibling directory
            };
            
            for (const auto& path : candidates) {
                NSURL* lib_url = [NSURL fileURLWithPath:[NSString stringWithUTF8String:path.c_str()]];
                
                // key: use [NSURL checkResourceIsReachableAndReturnError:] to avoid noise
                if ([lib_url checkResourceIsReachableAndReturnError:nil]) {
                    coreLib = [device newLibraryWithURL:lib_url error:&error];
                    if (coreLib) {
                         printf("metalcore: Loaded kernels from %s\n", path.c_str());
                         break;
                    }
                }
            }
        }
        
        // 2. Fallback: Hardcoded dev path (for local debugging only)
        if (!coreLib) {
             const char* dev_path = "/Users/kris/localprojects/metalops/packages/metalcore/src/metalcore/native/core_kernels.metallib";
             NSURL* lib_url = [NSURL fileURLWithPath:[NSString stringWithUTF8String:dev_path]];
             if ([lib_url checkResourceIsReachableAndReturnError:nil]) {
                  coreLib = [device newLibraryWithURL:lib_url error:&error];
                  if (coreLib) printf("metalcore: Loaded kernels from DEV path\n");
             }
        }

        // 3. Fallback: Compile from source (Developer mode / No .metallib found)
        if (!coreLib) {
             // Look for .metal source relative to dylib as well
             // ... implementation omitted for brevity, usually rely on precompiled in prod ...
             printf("metalcore: WARNING - Could not find .metallib. Falling back to source (if available).\n");
             
             // Try dev source path
             const char* src_path = "/Users/kris/localprojects/metalops/packages/metalcore/src/metalcore/native/core_kernels.metal";
             std::ifstream file(src_path);
            
            if (file.good()) {
                std::stringstream buffer;
                buffer << file.rdbuf();
                std::string content = buffer.str();
                
                NSString* src = [NSString stringWithUTF8String:content.c_str()];
                MTLCompileOptions* opts = [[MTLCompileOptions alloc] init];
                opts.mathMode = MTLMathModeFast;
                
                coreLib = [device newLibraryWithSource:src options:opts error:&error];
                
                if (coreLib) {
                    printf("metalcore: Compiled kernels from DEV source %s\n", src_path);
                } else {
                    printf("metalcore: Failed to compile source: %s\n", [[error localizedDescription] UTF8String]);
                    return;
                }
            } else {
                printf("metalcore: FATAL - Could not find .metallib or .metal kernel source.\n");
                return;
            }
        }
        
        // Load functions
        kernels.geqr2 = [coreLib newFunctionWithName:@"geqr2_panel_kernel"];
        kernels.geqr2Fused = [coreLib newFunctionWithName:@"geqr2_fused_kernel"];
        kernels.householder = [coreLib newFunctionWithName:@"householder_vector_kernel"];
        kernels.applyHouseholder = [coreLib newFunctionWithName:@"apply_householder_kernel"];
        kernels.larfbStep1 = [coreLib newFunctionWithName:@"larfb_step1_vtc"];
        kernels.larfbStep2 = [coreLib newFunctionWithName:@"larfb_step2_tw"];
        kernels.larfbStep3 = [coreLib newFunctionWithName:@"larfb_step3_cvw"];
        kernels.larft = [coreLib newFunctionWithName:@"larft_kernel"];
        kernels.trsmLower = [coreLib newFunctionWithName:@"trsm_lower_kernel"];
        kernels.trsmUpper = [coreLib newFunctionWithName:@"trsm_upper_kernel"];
        kernels.qrFullFused = [coreLib newFunctionWithName:@"qr_full_fused_kernel"];
        kernels.qrBatched = [coreLib newFunctionWithName:@"qr_batched_kernel"];
        kernels.trsmBatched = [coreLib newFunctionWithName:@"trsm_batched_kernel"];
        kernels.columnNorms = [coreLib newFunctionWithName:@"column_norms_kernel"];
        kernels.choleskyBatched = [coreLib newFunctionWithName:@"cholesky_batched_kernel"];
        kernels.trsmLowerBatched = [coreLib newFunctionWithName:@"trsm_lower_batched_kernel"];
        kernels.trsmUpperBatched = [coreLib newFunctionWithName:@"trsm_upper_batched_kernel"];
        kernels.choleskySolveBatched = [coreLib newFunctionWithName:@"cholesky_solve_batched_kernel"];
        
        // New optimization kernels
        kernels.columnNormSort = [coreLib newFunctionWithName:@"column_norm_sort_kernel"];
        kernels.signCanonicalize = [coreLib newFunctionWithName:@"sign_canonicalize_kernel"];
        kernels.batchedQtB = [coreLib newFunctionWithName:@"batched_qt_b_kernel"];
        
        // High-impact ML/LA kernels
        kernels.luBatched = [coreLib newFunctionWithName:@"lu_batched_kernel"];
        kernels.inverseBatched = [coreLib newFunctionWithName:@"inverse_batched_kernel"];
        kernels.syrkBatched = [coreLib newFunctionWithName:@"syrk_batched_kernel"];
        kernels.frobeniusNormBatched = [coreLib newFunctionWithName:@"frobenius_norm_batched_kernel"];
        kernels.softmaxBatched = [coreLib newFunctionWithName:@"softmax_batched_kernel"];
        kernels.traceBatched = [coreLib newFunctionWithName:@"trace_batched_kernel"];
        
        // Create pipeline states
        if (kernels.geqr2) {
            kernels.geqr2PSO = [device newComputePipelineStateWithFunction:kernels.geqr2 error:&error];
            if (!kernels.geqr2PSO) printf("Failed to create geqr2PSO: %s\n", [[error localizedDescription] UTF8String]);
        }
        if (kernels.geqr2Fused) {
            kernels.geqr2FusedPSO = [device newComputePipelineStateWithFunction:kernels.geqr2Fused error:&error];
            if (!kernels.geqr2FusedPSO) printf("Failed to create geqr2FusedPSO: %s\n", [[error localizedDescription] UTF8String]);
        }
        if (kernels.householder) {
            kernels.householderPSO = [device newComputePipelineStateWithFunction:kernels.householder error:&error];
        }
        if (kernels.applyHouseholder) {
            kernels.applyHouseholderPSO = [device newComputePipelineStateWithFunction:kernels.applyHouseholder error:&error];
        }
        if (kernels.larfbStep1) {
            kernels.larfbStep1PSO = [device newComputePipelineStateWithFunction:kernels.larfbStep1 error:&error];
        }
        if (kernels.larfbStep2) {
            kernels.larfbStep2PSO = [device newComputePipelineStateWithFunction:kernels.larfbStep2 error:&error];
        }
        if (kernels.larfbStep3) {
            kernels.larfbStep3PSO = [device newComputePipelineStateWithFunction:kernels.larfbStep3 error:&error];
        }
        if (kernels.larft) {
            kernels.larftPSO = [device newComputePipelineStateWithFunction:kernels.larft error:&error];
        }
        if (kernels.trsmLower) {
            kernels.trsmLowerPSO = [device newComputePipelineStateWithFunction:kernels.trsmLower error:&error];
        }
        if (kernels.trsmUpper) {
            kernels.trsmUpperPSO = [device newComputePipelineStateWithFunction:kernels.trsmUpper error:&error];
        }
        if (kernels.qrFullFused) {
            kernels.qrFullFusedPSO = [device newComputePipelineStateWithFunction:kernels.qrFullFused error:&error];
            if (!kernels.qrFullFusedPSO) printf("Failed to create qrFullFusedPSO: %s\n", [[error localizedDescription] UTF8String]);
        }
        if (kernels.qrBatched) {
            kernels.qrBatchedPSO = [device newComputePipelineStateWithFunction:kernels.qrBatched error:&error];
            if (!kernels.qrBatchedPSO) printf("Failed to create qrBatchedPSO: %s\n", [[error localizedDescription] UTF8String]);
        }
        if (kernels.trsmBatched) {
            kernels.trsmBatchedPSO = [device newComputePipelineStateWithFunction:kernels.trsmBatched error:&error];
            if (!kernels.trsmBatchedPSO) printf("Failed to create trsmBatchedPSO: %s\n", [[error localizedDescription] UTF8String]);
        }
        if (kernels.columnNorms) {
            kernels.columnNormsPSO = [device newComputePipelineStateWithFunction:kernels.columnNorms error:&error];
            if (!kernels.columnNormsPSO) printf("Failed to create columnNormsPSO: %s\n", [[error localizedDescription] UTF8String]);
        }
        if (kernels.choleskyBatched) {
            kernels.choleskyBatchedPSO = [device newComputePipelineStateWithFunction:kernels.choleskyBatched error:&error];
            if (!kernels.choleskyBatchedPSO) printf("Failed to create choleskyBatchedPSO: %s\n", [[error localizedDescription] UTF8String]);
        }
        if (kernels.trsmLowerBatched) {
            kernels.trsmLowerBatchedPSO = [device newComputePipelineStateWithFunction:kernels.trsmLowerBatched error:&error];
            if (!kernels.trsmLowerBatchedPSO) printf("Failed to create trsmLowerBatchedPSO: %s\n", [[error localizedDescription] UTF8String]);
        }
        if (kernels.trsmUpperBatched) {
            kernels.trsmUpperBatchedPSO = [device newComputePipelineStateWithFunction:kernels.trsmUpperBatched error:&error];
            if (!kernels.trsmUpperBatchedPSO) printf("Failed to create trsmUpperBatchedPSO: %s\n", [[error localizedDescription] UTF8String]);
        }
        if (kernels.choleskySolveBatched) {
            kernels.choleskySolveBatchedPSO = [device newComputePipelineStateWithFunction:kernels.choleskySolveBatched error:&error];
            if (!kernels.choleskySolveBatchedPSO) printf("Failed to create choleskySolveBatchedPSO: %s\n", [[error localizedDescription] UTF8String]);
        }
        
        // New optimization kernels
        if (kernels.columnNormSort) {
            kernels.columnNormSortPSO = [device newComputePipelineStateWithFunction:kernels.columnNormSort error:&error];
            if (!kernels.columnNormSortPSO) printf("Failed to create columnNormSortPSO: %s\n", [[error localizedDescription] UTF8String]);
        }
        if (kernels.signCanonicalize) {
            kernels.signCanonicalizePSO = [device newComputePipelineStateWithFunction:kernels.signCanonicalize error:&error];
            if (!kernels.signCanonicalizePSO) printf("Failed to create signCanonicalizePSO: %s\n", [[error localizedDescription] UTF8String]);
        }
        if (kernels.batchedQtB) {
            kernels.batchedQtBPSO = [device newComputePipelineStateWithFunction:kernels.batchedQtB error:&error];
            if (!kernels.batchedQtBPSO) printf("Failed to create batchedQtBPSO: %s\n", [[error localizedDescription] UTF8String]);
        }
        
        // High-impact ML/LA kernels
        if (kernels.luBatched) {
            kernels.luBatchedPSO = [device newComputePipelineStateWithFunction:kernels.luBatched error:&error];
            if (!kernels.luBatchedPSO) printf("Failed to create luBatchedPSO: %s\n", [[error localizedDescription] UTF8String]);
        }
        if (kernels.inverseBatched) {
            kernels.inverseBatchedPSO = [device newComputePipelineStateWithFunction:kernels.inverseBatched error:&error];
            if (!kernels.inverseBatchedPSO) printf("Failed to create inverseBatchedPSO: %s\n", [[error localizedDescription] UTF8String]);
        }
        if (kernels.syrkBatched) {
            kernels.syrkBatchedPSO = [device newComputePipelineStateWithFunction:kernels.syrkBatched error:&error];
            if (!kernels.syrkBatchedPSO) printf("Failed to create syrkBatchedPSO: %s\n", [[error localizedDescription] UTF8String]);
        }
        if (kernels.frobeniusNormBatched) {
            kernels.frobeniusNormBatchedPSO = [device newComputePipelineStateWithFunction:kernels.frobeniusNormBatched error:&error];
            if (!kernels.frobeniusNormBatchedPSO) printf("Failed to create frobeniusNormBatchedPSO: %s\n", [[error localizedDescription] UTF8String]);
        }
        if (kernels.softmaxBatched) {
            kernels.softmaxBatchedPSO = [device newComputePipelineStateWithFunction:kernels.softmaxBatched error:&error];
            if (!kernels.softmaxBatchedPSO) printf("Failed to create softmaxBatchedPSO: %s\n", [[error localizedDescription] UTF8String]);
        }
        if (kernels.traceBatched) {
            kernels.traceBatchedPSO = [device newComputePipelineStateWithFunction:kernels.traceBatched error:&error];
            if (!kernels.traceBatchedPSO) printf("Failed to create traceBatchedPSO: %s\n", [[error localizedDescription] UTF8String]);
        }
        
        printf("metalcore: Loaded %d kernel functions\n", 
            (kernels.geqr2 ? 1 : 0) + (kernels.householder ? 1 : 0) + 
            (kernels.applyHouseholder ? 1 : 0) + (kernels.larfbStep1 ? 1 : 0) +
            (kernels.larfbStep2 ? 1 : 0) + (kernels.larfbStep3 ? 1 : 0) +
            (kernels.larft ? 1 : 0) + (kernels.trsmLower ? 1 : 0) + 
            (kernels.trsmUpper ? 1 : 0) + (kernels.qrFullFused ? 1 : 0) +
            (kernels.qrBatched ? 1 : 0) + (kernels.trsmBatched ? 1 : 0) +
            (kernels.columnNorms ? 1 : 0) + (kernels.choleskyBatched ? 1 : 0) +
            (kernels.trsmLowerBatched ? 1 : 0) + (kernels.trsmUpperBatched ? 1 : 0));
    });
}

// -----------------------------------------------------------------------------
// Triangular Solve
// -----------------------------------------------------------------------------

torch::Tensor trsm_metal(
    torch::Tensor A,
    torch::Tensor b,
    bool lower,
    bool transpose
) {
    load_core_kernels();
    
    TORCH_CHECK(A.device().type() == at::kMPS, "A must be on MPS device");
    TORCH_CHECK(b.device().type() == at::kMPS, "b must be on MPS device");
    TORCH_CHECK(A.dim() == 2, "A must be 2D");
    TORCH_CHECK(A.size(0) == A.size(1), "A must be square");
    
    int64_t N = A.size(0);
    
    // Make contiguous
    auto A_contig = A.contiguous();
    auto x = b.clone().contiguous();
    
    id<MTLComputePipelineState> pso = lower ? kernels.trsmLowerPSO : kernels.trsmUpperPSO;
    
    if (!pso) {
        // Fallback to PyTorch - move to CPU, solve, move back
        auto A_cpu = A.cpu();
        auto x_cpu = b.cpu().clone();
        
        // Manual triangular solve on CPU
        int64_t n = A_cpu.size(0);
        auto A_acc = A_cpu.accessor<float, 2>();
        auto x_acc = x_cpu.accessor<float, 1>();
        
        if (lower) {
            for (int64_t i = 0; i < n; i++) {
                float sum = 0.0f;
                for (int64_t j = 0; j < i; j++) {
                    sum += A_acc[i][j] * x_acc[j];
                }
                x_acc[i] = (x_acc[i] - sum) / A_acc[i][i];
            }
        } else {
            for (int64_t i = n - 1; i >= 0; i--) {
                float sum = 0.0f;
                for (int64_t j = i + 1; j < n; j++) {
                    sum += A_acc[i][j] * x_acc[j];
                }
                x_acc[i] = (x_acc[i] - sum) / A_acc[i][i];
            }
        }
        return x_cpu.to(A.device());
    }
    
    @autoreleasepool {
        MPSStream* stream = getCurrentMPSStream();
        id<MTLComputeCommandEncoder> encoder = stream->commandEncoder();
        
        [encoder setComputePipelineState:pso];
        [encoder setBuffer:getMTLBufferStorage(A_contig) offset:A_contig.storage_offset() * A_contig.element_size() atIndex:0];
        [encoder setBuffer:getMTLBufferStorage(x) offset:x.storage_offset() * x.element_size() atIndex:1];
        
        uint32_t N_uint = (uint32_t)N;
        [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:2];
        
        [encoder dispatchThreads:MTLSizeMake(1, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
        
        stream->synchronize(SyncType::COMMIT_AND_WAIT);
    }
    
    return x;
}

// -----------------------------------------------------------------------------
// Panel QR (geqr2)
// -----------------------------------------------------------------------------

std::tuple<torch::Tensor, torch::Tensor> geqr2_metal(torch::Tensor A) {
    load_core_kernels();
    
    TORCH_CHECK(A.device().type() == at::kMPS, "A must be on MPS device");
    TORCH_CHECK(A.dim() == 2, "A must be 2D");
    
    int64_t M = A.size(0);
    int64_t N = A.size(1);
    int64_t K = std::min(M, N);
    
    // Output: R (with V stored below diagonal) and tau
    auto R = A.clone().contiguous();
    auto tau = torch::zeros({K}, A.options());
    
    if (!kernels.geqr2PSO) {
        // CPU fallback - manual Householder QR
        auto R_cpu = A.cpu().clone();
        auto tau_cpu = torch::zeros({K}, torch::dtype(torch::kFloat32));
        
        auto R_acc = R_cpu.accessor<float, 2>();
        auto tau_acc = tau_cpu.accessor<float, 1>();
        
        for (int64_t k = 0; k < K; k++) {
            // Compute norm of column k below diagonal
            float sigma = 0.0f;
            for (int64_t i = k + 1; i < M; i++) {
                sigma += R_acc[i][k] * R_acc[i][k];
            }
            
            float x0 = R_acc[k][k];
            if (sigma < 1e-10f) {
                tau_acc[k] = 0.0f;
                continue;
            }
            
            float norm_x = std::sqrt(x0 * x0 + sigma);
            float sign = (x0 >= 0.0f) ? 1.0f : -1.0f;
            float v0 = x0 + sign * norm_x;
            float tau_k = 2.0f * v0 * v0 / (v0 * v0 + sigma);
            tau_acc[k] = tau_k;
            
            // Update diagonal
            R_acc[k][k] = -sign * norm_x;
            
            // Apply Householder to trailing columns
            for (int64_t j = k + 1; j < N; j++) {
                // Compute v^T @ A[k:, j]
                float dot = R_acc[k][j];  // v[0] = 1
                for (int64_t i = k + 1; i < M; i++) {
                    float v_i = R_acc[i][k] / v0;
                    dot += v_i * R_acc[i][j];
                }
                
                // Update column
                R_acc[k][j] -= tau_k * dot;
                for (int64_t i = k + 1; i < M; i++) {
                    float v_i = R_acc[i][k] / v0;
                    R_acc[i][j] -= tau_k * v_i * dot;
                }
            }
            
            // Store v below diagonal (normalized)
            for (int64_t i = k + 1; i < M; i++) {
                R_acc[i][k] /= v0;
            }
        }
        
        return std::make_tuple(R_cpu.to(A.device()), tau_cpu.to(A.device()));
    }
    
    // Use fused kernel if available and panel fits in shared memory
    // Shared memory: M*N floats for panel + 256 floats for reduction buffer
    bool use_fused = kernels.geqr2FusedPSO && (M * N + 256) * sizeof(float) <= 32768;
    
    @autoreleasepool {
        MPSStream* stream = getCurrentMPSStream();
        id<MTLComputeCommandEncoder> encoder = stream->commandEncoder();
        
        if (use_fused) {
            [encoder setComputePipelineState:kernels.geqr2FusedPSO];
            [encoder setBuffer:getMTLBufferStorage(R) offset:R.storage_offset() * R.element_size() atIndex:0];
            [encoder setBuffer:getMTLBufferStorage(tau) offset:0 atIndex:1];
            
            uint32_t M_uint = (uint32_t)M;
            uint32_t N_uint = (uint32_t)N;
            uint32_t lda = (uint32_t)N;  // Row-major: lda = N
            [encoder setBytes:&M_uint length:sizeof(M_uint) atIndex:2];
            [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:3];
            [encoder setBytes:&lda length:sizeof(lda) atIndex:4];
            
            // Allocate shared memory for panel + reduction buffer
            NSUInteger shared_size = (M * N + 256) * sizeof(float);
            [encoder setThreadgroupMemoryLength:shared_size atIndex:0];
            
            // Launch single threadgroup with 256 threads
            NSUInteger tg_size = 256;
            [encoder dispatchThreadgroups:MTLSizeMake(1, 1, 1) 
                threadsPerThreadgroup:MTLSizeMake(tg_size, 1, 1)];
        } else {
            // Fallback to original kernel
            [encoder setComputePipelineState:kernels.geqr2PSO];
            [encoder setBuffer:getMTLBufferStorage(R) offset:R.storage_offset() * R.element_size() atIndex:0];
            [encoder setBuffer:getMTLBufferStorage(tau) offset:0 atIndex:1];
            
            uint32_t M_uint = (uint32_t)M;
            uint32_t N_uint = (uint32_t)N;
            uint32_t lda = (uint32_t)N;
            [encoder setBytes:&M_uint length:sizeof(M_uint) atIndex:2];
            [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:3];
            [encoder setBytes:&lda length:sizeof(lda) atIndex:4];
            
            NSUInteger tg_size = 256;
            NSUInteger shared_size = tg_size * sizeof(float);
            [encoder setThreadgroupMemoryLength:shared_size atIndex:0];
            [encoder dispatchThreads:MTLSizeMake(tg_size, 1, 1) 
                threadsPerThreadgroup:MTLSizeMake(tg_size, 1, 1)];
        }
        
        stream->synchronize(SyncType::COMMIT_AND_WAIT);
    }
    
    return std::make_tuple(R, tau);
}

// -----------------------------------------------------------------------------
// Block Reflector Application (larfb)
// Apply H = I - V @ T @ V^T to C
// C = C - V @ (T @ (V^T @ C)) for trans=false
// C = C - V @ (T^T @ (V^T @ C)) for trans=true
// -----------------------------------------------------------------------------

torch::Tensor larfb_metal(
    torch::Tensor C,
    torch::Tensor V,
    torch::Tensor T,
    bool trans,
    int64_t panel_start
) {
    load_core_kernels();
    
    TORCH_CHECK(C.device().type() == at::kMPS, "C must be on MPS device");
    TORCH_CHECK(V.device().type() == at::kMPS, "V must be on MPS device");
    TORCH_CHECK(T.device().type() == at::kMPS, "T must be on MPS device");
    
    int64_t M = C.size(0);
    int64_t N = C.size(1);
    int64_t K = V.size(1);  // Number of reflectors
    
    // Build V_full efficiently using tensor operations
    // V_full has: 1s on diagonal (at panel_start offset), 0s above, V below
    torch::Tensor V_full;
    
    int64_t V_rows = V.size(0);
    int64_t V_cols = V.size(1);
    
    if (panel_start == 0 && V_rows >= V_cols) {
        // Fast path: use tril to zero above diagonal, then set diagonal to 1
        // tril keeps diagonal and below
        V_full = torch::tril(V, -1);  // Keep below diagonal only (strict lower)
        // Add identity for the diagonal
        auto diag_ones = torch::eye(V_rows, V_cols, V.options());
        V_full = V_full + diag_ones;
    } else {
        // General case - still need element-wise but less common
        V_full = V.clone();
        for (int64_t k = 0; k < K; k++) {
            if (panel_start + k < M) {
                V_full.index_put_({panel_start + k, k}, 1.0f);
            }
            for (int64_t i = 0; i < panel_start + k; i++) {
                V_full.index_put_({i, k}, 0.0f);
            }
        }
    }
    
    // Use optimized mm (faster than matmul for 2D tensors)
    // W = V^T @ C  (K x N)
    auto W = V_full.t().mm(C);
    
    // W = T^T @ W or T @ W  (K x N)
    if (trans) {
        W = T.t().mm(W);
    } else {
        W = T.mm(W);
    }
    
    // C = C - V @ W  (M x N) - use sub_ variant for efficiency
    return C - V_full.mm(W);
}

// -----------------------------------------------------------------------------
// Build T matrix (larft)
// -----------------------------------------------------------------------------

torch::Tensor larft_metal(
    torch::Tensor V,
    torch::Tensor tau,
    int64_t panel_start
) {
    load_core_kernels();
    
    TORCH_CHECK(V.device().type() == at::kMPS, "V must be on MPS device");
    TORCH_CHECK(tau.device().type() == at::kMPS, "tau must be on MPS device");
    
    int64_t M = V.size(0);
    int64_t K = V.size(1);
    
    // Always use CPU fallback for larft - it's small and sequential
    auto T_cpu = torch::zeros({K, K}, torch::dtype(torch::kFloat32));
    auto V_cpu = V.cpu();
    auto tau_cpu = tau.cpu();
    
    auto T_acc = T_cpu.accessor<float, 2>();
    auto V_acc = V_cpu.accessor<float, 2>();
    auto tau_acc = tau_cpu.accessor<float, 1>();
    
    // Build T column by column using LAPACK algorithm:
    // T[i,i] = tau[i]
    // T[0:i, i] = -tau[i] * T[0:i, 0:i] @ V[:, 0:i]^T @ V[:, i]
    for (int64_t i = 0; i < K; i++) {
        T_acc[i][i] = tau_acc[i];
        
        if (i > 0) {
            // Step 1: Compute w = V[:, 0:i]^T @ V[:, i]
            // w[j] = V[:, j]^T @ V[:, i] for j = 0..i-1
            std::vector<float> w(i);
            for (int64_t j = 0; j < i; j++) {
                float dot = 0.0f;
                for (int64_t m = 0; m < M; m++) {
                    float vj, vi;
                    
                    // V[m, j] with implicit 1 at (panel_start + j)
                    if (m == panel_start + j) vj = 1.0f;
                    else if (m < panel_start + j) vj = 0.0f;
                    else vj = V_acc[m][j];
                    
                    // V[m, i] with implicit 1 at (panel_start + i)
                    if (m == panel_start + i) vi = 1.0f;
                    else if (m < panel_start + i) vi = 0.0f;
                    else vi = V_acc[m][i];
                    
                    dot += vj * vi;
                }
                w[j] = dot;
            }
            
            // Step 2: T[0:i, i] = -tau[i] * T[0:i, 0:i] @ w
            // This is a triangular matrix-vector product (T is upper triangular)
            for (int64_t j = 0; j < i; j++) {
                float sum = 0.0f;
                for (int64_t k = j; k < i; k++) {  // T[j,k] for k >= j (upper triangular)
                    sum += T_acc[j][k] * w[k];
                }
                T_acc[j][i] = -tau_acc[i] * sum;
            }
        }
    }
    
    return T_cpu.to(V.device());
}

// -----------------------------------------------------------------------------
// Full Blocked QR
// -----------------------------------------------------------------------------

std::tuple<torch::Tensor, torch::Tensor> qr_metal(torch::Tensor A, int64_t block_size) {
    load_core_kernels();
    
    TORCH_CHECK(A.device().type() == at::kMPS, "A must be on MPS device");
    TORCH_CHECK(A.dim() == 2, "A must be 2D");
    
    int64_t M = A.size(0);
    int64_t N = A.size(1);
    int64_t K = std::min(M, N);
    
    if (block_size <= 0) block_size = 32;
    
    // Copy A to R (will be modified in place)
    auto R = A.clone();
    
    // Storage for building Q  
    std::vector<std::tuple<int64_t, torch::Tensor, torch::Tensor>> panels;
    panels.reserve((K + block_size - 1) / block_size);
    
    for (int64_t j = 0; j < K; j += block_size) {
        int64_t jb = std::min(block_size, K - j);
        
        // Extract panel R[j:, j:j+jb] - need contiguous for kernel
        auto panel = R.index({torch::indexing::Slice(j, M), torch::indexing::Slice(j, j + jb)}).contiguous();
        
        // Factor panel: panel -> R_panel with V below diagonal, tau
        auto [R_panel, tau_panel] = geqr2_metal(panel);
        
        // Copy R_panel back to R
        R.index_put_({torch::indexing::Slice(j, M), torch::indexing::Slice(j, j + jb)}, R_panel);
        
        // V_panel is stored in R below the diagonal - reuse R_panel
        // Build T matrix for this panel
        auto T_panel = larft_metal(R_panel, tau_panel, 0);
        
        // Apply block reflector to trailing matrix R[j:, j+jb:]
        if (j + jb < N) {
            auto trailing = R.index({torch::indexing::Slice(j, M), torch::indexing::Slice(j + jb, N)}).contiguous();
            auto trailing_updated = larfb_metal(trailing, R_panel, T_panel, true, 0);
            R.index_put_({torch::indexing::Slice(j, M), torch::indexing::Slice(j + jb, N)}, trailing_updated);
        }
        
        panels.push_back(std::make_tuple(j, R_panel, T_panel));
    }
    
    // Zero below diagonal of R efficiently using triu
    auto R_upper = torch::triu(R.index({torch::indexing::Slice(0, K), torch::indexing::Slice()}));
    
    // Build Q by applying reflectors in reverse
    auto Q = torch::eye(M, K, A.options());
    
    for (auto it = panels.rbegin(); it != panels.rend(); ++it) {
        auto& [j, V_panel, T_panel] = *it;
        
        auto Q_sub = Q.index({torch::indexing::Slice(j, M), torch::indexing::Slice(j, K)}).contiguous();
        auto Q_updated = larfb_metal(Q_sub, V_panel, T_panel, false, 0);
        Q.index_put_({torch::indexing::Slice(j, M), torch::indexing::Slice(j, K)}, Q_updated);
    }
    
    return std::make_tuple(Q, R_upper);
}

// -----------------------------------------------------------------------------
// Fully Fused QR - Single Metal Dispatch
// -----------------------------------------------------------------------------

std::tuple<torch::Tensor, torch::Tensor> qr_fused_metal(torch::Tensor A) {
    load_core_kernels();
    
    TORCH_CHECK(A.device().type() == at::kMPS, "A must be on MPS device");
    TORCH_CHECK(A.dim() == 2, "A must be 2D");
    
    int64_t M = A.size(0);
    int64_t N = A.size(1);
    int64_t K = std::min(M, N);
    
    // Check if matrix fits in shared memory
    // Shared memory layout: M*N (R) + M*K (Q) + K (tau) + 256 (reduction) floats
    // 32KB = 8192 floats
    int64_t shared_needed = M * N + M * K + K + 256;
    
    if (!kernels.qrFullFusedPSO || shared_needed > 8000) {
        // Fall back to blocked QR for large matrices
        return qr_metal(A, 32);
    }
    
    // Ensure input is contiguous and row-major
    auto A_in = A.contiguous();
    
    // Allocate outputs
    auto Q_out = torch::zeros({M, K}, A.options());
    auto R_out = torch::zeros({K, N}, A.options());
    
    @autoreleasepool {
        MPSStream* stream = getCurrentMPSStream();
        id<MTLComputeCommandEncoder> encoder = stream->commandEncoder();
        
        [encoder setComputePipelineState:kernels.qrFullFusedPSO];
        [encoder setBuffer:getMTLBufferStorage(A_in) offset:A_in.storage_offset() * A_in.element_size() atIndex:0];
        [encoder setBuffer:getMTLBufferStorage(Q_out) offset:0 atIndex:1];
        [encoder setBuffer:getMTLBufferStorage(R_out) offset:0 atIndex:2];
        
        uint32_t M_uint = (uint32_t)M;
        uint32_t N_uint = (uint32_t)N;
        [encoder setBytes:&M_uint length:sizeof(M_uint) atIndex:3];
        [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:4];
        
        // Allocate shared memory
        NSUInteger shared_size = shared_needed * sizeof(float);
        [encoder setThreadgroupMemoryLength:shared_size atIndex:0];
        
        // Launch single threadgroup with 256 threads
        NSUInteger tg_size = 256;
        [encoder dispatchThreadgroups:MTLSizeMake(1, 1, 1) 
            threadsPerThreadgroup:MTLSizeMake(tg_size, 1, 1)];
        
        stream->synchronize(SyncType::COMMIT_AND_WAIT);
    }
    
    return std::make_tuple(Q_out, R_out);
}

// -----------------------------------------------------------------------------
// Batched QR - Process multiple matrices in single dispatch
// -----------------------------------------------------------------------------

std::tuple<torch::Tensor, torch::Tensor> qr_batched_metal(torch::Tensor A_batch) {
    load_core_kernels();
    
    TORCH_CHECK(A_batch.device().type() == at::kMPS, "A must be on MPS device");
    TORCH_CHECK(A_batch.dim() == 3, "A must be 3D (Batch, M, N)");
    
    int64_t Batch = A_batch.size(0);
    int64_t M = A_batch.size(1);
    int64_t N = A_batch.size(2);
    int64_t K = std::min(M, N);
    
    // Check if single matrix fits in shared memory
    int64_t shared_per_matrix = M * N + M * K + K + 256;
    
    if (!kernels.qrBatchedPSO || shared_per_matrix > 8000) {
        // Fall back to sequential processing
        auto Q_list = torch::zeros({Batch, M, K}, A_batch.options());
        auto R_list = torch::zeros({Batch, K, N}, A_batch.options());
        
        for (int64_t b = 0; b < Batch; b++) {
            auto [Q, R] = qr_fused_metal(A_batch[b]);
            Q_list[b] = Q;
            R_list[b] = R;
        }
        return std::make_tuple(Q_list, R_list);
    }
    
    auto A_contig = A_batch.contiguous();
    auto Q_out = torch::zeros({Batch, M, K}, A_batch.options());
    auto R_out = torch::zeros({Batch, K, N}, A_batch.options());
    
    @autoreleasepool {
        MPSStream* stream = getCurrentMPSStream();
        id<MTLComputeCommandEncoder> encoder = stream->commandEncoder();
        
        [encoder setComputePipelineState:kernels.qrBatchedPSO];
        [encoder setBuffer:getMTLBufferStorage(A_contig) offset:A_contig.storage_offset() * A_contig.element_size() atIndex:0];
        [encoder setBuffer:getMTLBufferStorage(Q_out) offset:0 atIndex:1];
        [encoder setBuffer:getMTLBufferStorage(R_out) offset:0 atIndex:2];
        
        uint32_t M_uint = (uint32_t)M;
        uint32_t N_uint = (uint32_t)N;
        uint32_t Batch_uint = (uint32_t)Batch;
        [encoder setBytes:&M_uint length:sizeof(M_uint) atIndex:3];
        [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:4];
        [encoder setBytes:&Batch_uint length:sizeof(Batch_uint) atIndex:5];
        
        NSUInteger shared_size = shared_per_matrix * sizeof(float);
        [encoder setThreadgroupMemoryLength:shared_size atIndex:0];
        
        // Launch Batch threadgroups, each with 256 threads
        NSUInteger tg_size = 256;
        [encoder dispatchThreadgroups:MTLSizeMake(Batch, 1, 1) 
            threadsPerThreadgroup:MTLSizeMake(tg_size, 1, 1)];
        
        stream->synchronize(SyncType::COMMIT_AND_WAIT);
    }
    
    return std::make_tuple(Q_out, R_out);
}

// -----------------------------------------------------------------------------
// Batched TRSM (Triangular Solve)
// -----------------------------------------------------------------------------

torch::Tensor trsm_batched_metal(torch::Tensor R, torch::Tensor B) {
    load_core_kernels();
    
    TORCH_CHECK(R.device().type() == at::kMPS, "R must be on MPS device");
    TORCH_CHECK(B.device().type() == at::kMPS, "B must be on MPS device");
    TORCH_CHECK(R.dim() == 3, "R must be 3D (Batch, N, N)");
    TORCH_CHECK(B.dim() == 3, "B must be 3D (Batch, N, NRHS)");
    
    int64_t Batch = R.size(0);
    int64_t N = R.size(1);
    int64_t NRHS = B.size(2);
    
    R = R.contiguous();
    B = B.contiguous();
    
    auto X = torch::empty({Batch, N, NRHS}, B.options());
    
    if (!kernels.trsmBatchedPSO) {
        // Fallback: sequential processing
        auto X_list = std::vector<torch::Tensor>();
        for (int64_t i = 0; i < Batch; i++) {
            auto R_i = R[i];
            auto B_i = B[i];
            // Simple back-substitution in C++
            auto X_i = torch::zeros_like(B_i);
            for (int64_t j = N - 1; j >= 0; j--) {
                auto sum = B_i.index({j}).clone();
                for (int64_t k = j + 1; k < N; k++) {
                    sum = sum - R_i.index({j, k}) * X_i.index({k});
                }
                X_i.index_put_({j}, sum / R_i.index({j, j}));
            }
            X_list.push_back(X_i);
        }
        return torch::stack(X_list);
    }
    
    @autoreleasepool {
        auto stream = at::mps::getCurrentMPSStream();
        auto cmdBuffer = stream->commandBuffer();
        auto encoder = [cmdBuffer computeCommandEncoder];
        
        [encoder setComputePipelineState:kernels.trsmBatchedPSO];
        [encoder setBuffer:getMTLBufferStorage(R) offset:R.storage_offset() * R.element_size() atIndex:0];
        [encoder setBuffer:getMTLBufferStorage(B) offset:B.storage_offset() * B.element_size() atIndex:1];
        [encoder setBuffer:getMTLBufferStorage(X) offset:0 atIndex:2];
        
        uint32_t N_uint = (uint32_t)N;
        uint32_t NRHS_uint = (uint32_t)NRHS;
        uint32_t Batch_uint = (uint32_t)Batch;
        [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:3];
        [encoder setBytes:&NRHS_uint length:sizeof(NRHS_uint) atIndex:4];
        [encoder setBytes:&Batch_uint length:sizeof(Batch_uint) atIndex:5];
        
        NSUInteger shared_size = N * NRHS * sizeof(float);
        [encoder setThreadgroupMemoryLength:shared_size atIndex:0];
        
        NSUInteger tg_size = 256;
        [encoder dispatchThreadgroups:MTLSizeMake(Batch, 1, 1) 
            threadsPerThreadgroup:MTLSizeMake(tg_size, 1, 1)];
        
        [encoder endEncoding];
        stream->synchronize(SyncType::COMMIT_AND_WAIT);
    }
    
    return X;
}

// -----------------------------------------------------------------------------
// Fused Solve Batched (QR + Q.T@b + TRSM in single command buffer)
// -----------------------------------------------------------------------------

torch::Tensor solve_batched_metal(torch::Tensor A, torch::Tensor b) {
    // Fused solve: eliminates Python overhead between QR, bmm, TRSM
    // All operations in single command buffer with one sync at end
    load_core_kernels();
    
    TORCH_CHECK(A.device().type() == at::kMPS, "A must be on MPS device");
    TORCH_CHECK(b.device().type() == at::kMPS, "b must be on MPS device");
    TORCH_CHECK(A.dim() == 3, "A must be 3D (Batch, N, N)");
    TORCH_CHECK(b.dim() == 3, "b must be 3D (Batch, N, K)");
    
    int64_t Batch = A.size(0);
    int64_t N = A.size(1);
    int64_t K = b.size(2);
    
    TORCH_CHECK(A.size(2) == N, "A must be square");
    TORCH_CHECK(b.size(1) == N, "b dimension mismatch");
    
    auto A_contig = A.contiguous();
    auto b_contig = b.contiguous();
    
    // Allocate outputs
    auto Q = torch::zeros({Batch, N, N}, A.options());
    auto R = torch::zeros({Batch, N, N}, A.options());
    auto c = torch::zeros({Batch, N, K}, b.options());  // Q.T @ b
    auto x = torch::zeros({Batch, N, K}, b.options());  // solution
    
    // Check kernel availability
    if (!kernels.qrBatchedPSO || !kernels.trsmBatchedPSO) {
        // Fallback: use existing separate functions
        auto [Q_out, R_out] = qr_batched_metal(A);
        auto c_out = torch::bmm(Q_out.transpose(-2, -1), b);
        return trsm_batched_metal(R_out, c_out);
    }
    
    @autoreleasepool {
        auto stream = at::mps::getCurrentMPSStream();
        auto cmdBuffer = stream->commandBuffer();
        
        // === Phase 1: Batched QR ===
        {
            id<MTLComputeCommandEncoder> encoder = [cmdBuffer computeCommandEncoder];
            
            [encoder setComputePipelineState:kernels.qrBatchedPSO];
            [encoder setBuffer:getMTLBufferStorage(A_contig) offset:A_contig.storage_offset() * A_contig.element_size() atIndex:0];
            [encoder setBuffer:getMTLBufferStorage(Q) offset:0 atIndex:1];
            [encoder setBuffer:getMTLBufferStorage(R) offset:0 atIndex:2];
            
            uint32_t N_uint = (uint32_t)N;
            uint32_t Batch_uint = (uint32_t)Batch;
            [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:3];
            [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:4];  // M = N for square
            [encoder setBytes:&Batch_uint length:sizeof(Batch_uint) atIndex:5];
            
            int64_t shared_per_matrix = N * N + N * N + N + 256;
            NSUInteger shared_size = shared_per_matrix * sizeof(float);
            [encoder setThreadgroupMemoryLength:shared_size atIndex:0];
            
            [encoder dispatchThreadgroups:MTLSizeMake(Batch, 1, 1) 
                threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
            
            [encoder endEncoding];
        }
        
        // === Phase 2: c = Q.T @ b (use custom Metal kernel - no sync needed!) ===
        if (kernels.batchedQtBPSO) {
            id<MTLComputeCommandEncoder> encoder = [cmdBuffer computeCommandEncoder];
            
            [encoder setComputePipelineState:kernels.batchedQtBPSO];
            [encoder setBuffer:getMTLBufferStorage(Q) offset:0 atIndex:0];
            [encoder setBuffer:getMTLBufferStorage(b_contig) offset:b_contig.storage_offset() * b_contig.element_size() atIndex:1];
            [encoder setBuffer:getMTLBufferStorage(c) offset:0 atIndex:2];
            
            uint32_t N_uint = (uint32_t)N;
            uint32_t K_uint = (uint32_t)K;
            uint32_t Batch_uint = (uint32_t)Batch;
            [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:3];  // M = N
            [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:4];  // N = N
            [encoder setBytes:&K_uint length:sizeof(K_uint) atIndex:5];
            [encoder setBytes:&Batch_uint length:sizeof(Batch_uint) atIndex:6];
            
            [encoder dispatchThreads:MTLSizeMake(K, N, Batch) 
                threadsPerThreadgroup:MTLSizeMake(16, 16, 1)];
            
            [encoder endEncoding];
        } else {
            // Fallback to torch::bmm (requires sync)
            stream->synchronize(SyncType::COMMIT_AND_WAIT);
            c = torch::bmm(Q.transpose(-2, -1), b_contig);
            cmdBuffer = stream->commandBuffer();  // Get new buffer after sync
        }
        
        // === Phase 3: Batched TRSM ===
        {
            c = c.contiguous();
            id<MTLComputeCommandEncoder> encoder = [cmdBuffer computeCommandEncoder];
            
            [encoder setComputePipelineState:kernels.trsmBatchedPSO];
            [encoder setBuffer:getMTLBufferStorage(R) offset:0 atIndex:0];
            [encoder setBuffer:getMTLBufferStorage(c) offset:0 atIndex:1];
            [encoder setBuffer:getMTLBufferStorage(x) offset:0 atIndex:2];
            
            uint32_t N_uint = (uint32_t)N;
            uint32_t K_uint = (uint32_t)K;
            uint32_t Batch_uint = (uint32_t)Batch;
            [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:3];
            [encoder setBytes:&K_uint length:sizeof(K_uint) atIndex:4];
            [encoder setBytes:&Batch_uint length:sizeof(Batch_uint) atIndex:5];
            
            NSUInteger shared_size = N * K * sizeof(float);
            [encoder setThreadgroupMemoryLength:shared_size atIndex:0];
            
            [encoder dispatchThreadgroups:MTLSizeMake(Batch, 1, 1) 
                threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
            
            [encoder endEncoding];
        }
        
        stream->synchronize(SyncType::COMMIT_AND_WAIT);
    }
    
    return x;
}

// -----------------------------------------------------------------------------
// Column Norm Sort (De Rijk optimization for SVD)
// -----------------------------------------------------------------------------

std::tuple<torch::Tensor, torch::Tensor> column_norm_sort_metal(torch::Tensor A) {
    load_core_kernels();
    
    TORCH_CHECK(A.device().type() == at::kMPS, "A must be on MPS device");
    TORCH_CHECK(A.dim() == 3, "A must be 3D (Batch, M, N)");
    
    int64_t Batch = A.size(0);
    int64_t M = A.size(1);
    int64_t N = A.size(2);
    
    auto A_contig = A.contiguous();
    auto A_sorted = torch::zeros_like(A);
    auto perm = torch::zeros({Batch, N}, torch::TensorOptions().dtype(torch::kInt32).device(A.device()));
    
    if (!kernels.columnNormSortPSO) {
        // Fallback: use Python-level sorting
        auto norms = torch::linalg_vector_norm(A, 2, /*dim=*/1);  // (B, N)
        auto argsort_result = torch::argsort(norms, /*dim=*/-1, /*descending=*/true);
        auto perm_out = argsort_result.to(torch::kInt32);
        auto A_out = torch::gather(A, 2, argsort_result.unsqueeze(1).expand_as(A));
        return std::make_tuple(A_out, perm_out);
    }
    
    @autoreleasepool {
        auto stream = at::mps::getCurrentMPSStream();
        auto cmdBuffer = stream->commandBuffer();
        auto encoder = [cmdBuffer computeCommandEncoder];
        
        [encoder setComputePipelineState:kernels.columnNormSortPSO];
        [encoder setBuffer:getMTLBufferStorage(A_contig) offset:A_contig.storage_offset() * A_contig.element_size() atIndex:0];
        [encoder setBuffer:getMTLBufferStorage(A_sorted) offset:0 atIndex:1];
        [encoder setBuffer:getMTLBufferStorage(perm) offset:0 atIndex:2];
        
        uint32_t M_uint = (uint32_t)M;
        uint32_t N_uint = (uint32_t)N;
        uint32_t Batch_uint = (uint32_t)Batch;
        [encoder setBytes:&M_uint length:sizeof(M_uint) atIndex:3];
        [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:4];
        [encoder setBytes:&Batch_uint length:sizeof(Batch_uint) atIndex:5];
        
        // Shared memory: N floats for norms + N ints for indices
        NSUInteger shared_size = N * (sizeof(float) + sizeof(int));
        [encoder setThreadgroupMemoryLength:shared_size atIndex:0];
        
        [encoder dispatchThreadgroups:MTLSizeMake(Batch, 1, 1) 
            threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        
        [encoder endEncoding];
        stream->synchronize(SyncType::COMMIT_AND_WAIT);
    }
    
    return std::make_tuple(A_sorted, perm);
}

// -----------------------------------------------------------------------------
// Sign Canonicalization (SVD U/V sign normalization)
// -----------------------------------------------------------------------------

void sign_canonicalize_metal(torch::Tensor U, torch::Tensor V) {
    load_core_kernels();
    
    TORCH_CHECK(U.device().type() == at::kMPS, "U must be on MPS device");
    TORCH_CHECK(V.device().type() == at::kMPS, "V must be on MPS device");
    TORCH_CHECK(U.dim() == 3, "U must be 3D (Batch, M, N)");
    TORCH_CHECK(V.dim() == 3, "V must be 3D (Batch, N, N)");
    
    int64_t Batch = U.size(0);
    int64_t M = U.size(1);
    int64_t N = U.size(2);
    
    if (!kernels.signCanonicalizePSO) {
        // Fallback: Python-level sign canonicalization
        auto max_vals = std::get<0>(torch::max(torch::abs(U), 1));  // (B, N)
        auto max_signs = torch::sign(torch::gather(U, 1, std::get<1>(torch::max(torch::abs(U), 1)).unsqueeze(1)));
        // Skip fallback implementation for now - kernel should work
        return;
    }
    
    @autoreleasepool {
        auto stream = at::mps::getCurrentMPSStream();
        auto cmdBuffer = stream->commandBuffer();
        auto encoder = [cmdBuffer computeCommandEncoder];
        
        [encoder setComputePipelineState:kernels.signCanonicalizePSO];
        [encoder setBuffer:getMTLBufferStorage(U) offset:U.storage_offset() * U.element_size() atIndex:0];
        [encoder setBuffer:getMTLBufferStorage(V) offset:V.storage_offset() * V.element_size() atIndex:1];
        
        uint32_t M_uint = (uint32_t)M;
        uint32_t N_uint = (uint32_t)N;
        uint32_t Batch_uint = (uint32_t)Batch;
        [encoder setBytes:&M_uint length:sizeof(M_uint) atIndex:2];
        [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:3];
        [encoder setBytes:&Batch_uint length:sizeof(Batch_uint) atIndex:4];
        
        [encoder dispatchThreadgroups:MTLSizeMake(Batch, 1, 1) 
            threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        
        [encoder endEncoding];
        stream->synchronize(SyncType::COMMIT_AND_WAIT);
    }
}

// -----------------------------------------------------------------------------
// Batched Q.T @ b (for fused solve without sync)
// -----------------------------------------------------------------------------

torch::Tensor batched_qt_b_metal(torch::Tensor Q, torch::Tensor b) {
    load_core_kernels();
    
    TORCH_CHECK(Q.device().type() == at::kMPS, "Q must be on MPS device");
    TORCH_CHECK(b.device().type() == at::kMPS, "b must be on MPS device");
    TORCH_CHECK(Q.dim() == 3, "Q must be 3D (Batch, M, N)");
    TORCH_CHECK(b.dim() == 3, "b must be 3D (Batch, M, K)");
    
    int64_t Batch = Q.size(0);
    int64_t M = Q.size(1);
    int64_t N = Q.size(2);
    int64_t K = b.size(2);
    
    auto Q_contig = Q.contiguous();
    auto b_contig = b.contiguous();
    auto c = torch::zeros({Batch, N, K}, b.options());
    
    if (!kernels.batchedQtBPSO) {
        // Fallback: use PyTorch bmm
        return torch::bmm(Q.transpose(-2, -1), b);
    }
    
    @autoreleasepool {
        auto stream = at::mps::getCurrentMPSStream();
        auto cmdBuffer = stream->commandBuffer();
        auto encoder = [cmdBuffer computeCommandEncoder];
        
        [encoder setComputePipelineState:kernels.batchedQtBPSO];
        [encoder setBuffer:getMTLBufferStorage(Q_contig) offset:Q_contig.storage_offset() * Q_contig.element_size() atIndex:0];
        [encoder setBuffer:getMTLBufferStorage(b_contig) offset:b_contig.storage_offset() * b_contig.element_size() atIndex:1];
        [encoder setBuffer:getMTLBufferStorage(c) offset:0 atIndex:2];
        
        uint32_t M_uint = (uint32_t)M;
        uint32_t N_uint = (uint32_t)N;
        uint32_t K_uint = (uint32_t)K;
        uint32_t Batch_uint = (uint32_t)Batch;
        [encoder setBytes:&M_uint length:sizeof(M_uint) atIndex:3];
        [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:4];
        [encoder setBytes:&K_uint length:sizeof(K_uint) atIndex:5];
        [encoder setBytes:&Batch_uint length:sizeof(Batch_uint) atIndex:6];
        
        // Dispatch threads for each output element
        [encoder dispatchThreads:MTLSizeMake(K, N, Batch) 
            threadsPerThreadgroup:MTLSizeMake(16, 16, 1)];
        
        [encoder endEncoding];
        stream->synchronize(SyncType::COMMIT_AND_WAIT);
    }
    
    return c;
}

// -----------------------------------------------------------------------------
// Batched LU Decomposition
// -----------------------------------------------------------------------------

std::tuple<torch::Tensor, torch::Tensor> lu_batched_metal(torch::Tensor A) {
    load_core_kernels();
    
    TORCH_CHECK(A.device().type() == at::kMPS, "A must be on MPS device");
    TORCH_CHECK(A.dim() == 3, "A must be 3D (Batch, N, N)");
    TORCH_CHECK(A.size(1) == A.size(2), "A must be square");
    
    int64_t Batch = A.size(0);
    int64_t N = A.size(1);
    
    auto LU = A.clone().contiguous();
    auto pivots = torch::zeros({Batch, N}, torch::TensorOptions().dtype(torch::kInt32).device(A.device()));
    
    if (!kernels.luBatchedPSO) {
        // Fallback
        TORCH_CHECK(false, "LU kernel not available");
    }
    
    @autoreleasepool {
        auto stream = at::mps::getCurrentMPSStream();
        auto cmdBuffer = stream->commandBuffer();
        auto encoder = [cmdBuffer computeCommandEncoder];
        
        [encoder setComputePipelineState:kernels.luBatchedPSO];
        [encoder setBuffer:getMTLBufferStorage(LU) offset:0 atIndex:0];
        [encoder setBuffer:getMTLBufferStorage(pivots) offset:0 atIndex:1];
        
        uint32_t N_uint = (uint32_t)N;
        uint32_t Batch_uint = (uint32_t)Batch;
        [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:2];
        [encoder setBytes:&Batch_uint length:sizeof(Batch_uint) atIndex:3];
        
        [encoder dispatchThreadgroups:MTLSizeMake(Batch, 1, 1) 
            threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        
        [encoder endEncoding];
        stream->synchronize(SyncType::COMMIT_AND_WAIT);
    }
    
    return std::make_tuple(LU, pivots);
}

// -----------------------------------------------------------------------------
// Batched SYRK: C = A.T @ A
// -----------------------------------------------------------------------------

torch::Tensor syrk_batched_metal(torch::Tensor A) {
    load_core_kernels();
    
    TORCH_CHECK(A.device().type() == at::kMPS, "A must be on MPS device");
    TORCH_CHECK(A.dim() == 3, "A must be 3D (Batch, M, N)");
    
    int64_t Batch = A.size(0);
    int64_t M = A.size(1);
    int64_t N = A.size(2);
    
    auto A_contig = A.contiguous();
    auto C = torch::zeros({Batch, N, N}, A.options());
    
    if (!kernels.syrkBatchedPSO) {
        // Fallback to bmm
        return torch::bmm(A.transpose(-2, -1), A);
    }
    
    @autoreleasepool {
        auto stream = at::mps::getCurrentMPSStream();
        auto cmdBuffer = stream->commandBuffer();
        auto encoder = [cmdBuffer computeCommandEncoder];
        
        [encoder setComputePipelineState:kernels.syrkBatchedPSO];
        [encoder setBuffer:getMTLBufferStorage(A_contig) offset:A_contig.storage_offset() * A_contig.element_size() atIndex:0];
        [encoder setBuffer:getMTLBufferStorage(C) offset:0 atIndex:1];
        
        uint32_t M_uint = (uint32_t)M;
        uint32_t N_uint = (uint32_t)N;
        uint32_t Batch_uint = (uint32_t)Batch;
        [encoder setBytes:&M_uint length:sizeof(M_uint) atIndex:2];
        [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:3];
        [encoder setBytes:&Batch_uint length:sizeof(Batch_uint) atIndex:4];
        
        [encoder dispatchThreads:MTLSizeMake(N, N, Batch) 
            threadsPerThreadgroup:MTLSizeMake(16, 16, 1)];
        
        [encoder endEncoding];
        stream->synchronize(SyncType::COMMIT_AND_WAIT);
    }
    
    return C;
}

// -----------------------------------------------------------------------------
// Batched Frobenius Norm
// -----------------------------------------------------------------------------

torch::Tensor frobenius_norm_batched_metal(torch::Tensor A) {
    load_core_kernels();
    
    TORCH_CHECK(A.device().type() == at::kMPS, "A must be on MPS device");
    TORCH_CHECK(A.dim() == 3, "A must be 3D (Batch, M, N)");
    
    int64_t Batch = A.size(0);
    int64_t M = A.size(1);
    int64_t N = A.size(2);
    
    auto A_contig = A.contiguous();
    auto norms = torch::zeros({Batch}, A.options());
    
    if (!kernels.frobeniusNormBatchedPSO) {
        // Fallback
        return torch::linalg_matrix_norm(A, "fro");
    }
    
    @autoreleasepool {
        auto stream = at::mps::getCurrentMPSStream();
        auto cmdBuffer = stream->commandBuffer();
        auto encoder = [cmdBuffer computeCommandEncoder];
        
        [encoder setComputePipelineState:kernels.frobeniusNormBatchedPSO];
        [encoder setBuffer:getMTLBufferStorage(A_contig) offset:A_contig.storage_offset() * A_contig.element_size() atIndex:0];
        [encoder setBuffer:getMTLBufferStorage(norms) offset:0 atIndex:1];
        
        uint32_t M_uint = (uint32_t)M;
        uint32_t N_uint = (uint32_t)N;
        uint32_t Batch_uint = (uint32_t)Batch;
        [encoder setBytes:&M_uint length:sizeof(M_uint) atIndex:2];
        [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:3];
        [encoder setBytes:&Batch_uint length:sizeof(Batch_uint) atIndex:4];
        
        NSUInteger shared_size = 256 * sizeof(float);
        [encoder setThreadgroupMemoryLength:shared_size atIndex:0];
        
        [encoder dispatchThreadgroups:MTLSizeMake(Batch, 1, 1) 
            threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        
        [encoder endEncoding];
        stream->synchronize(SyncType::COMMIT_AND_WAIT);
    }
    
    return norms;
}

// -----------------------------------------------------------------------------
// Batched Softmax
// -----------------------------------------------------------------------------

torch::Tensor softmax_batched_metal(torch::Tensor x, float temperature) {
    load_core_kernels();
    
    TORCH_CHECK(x.device().type() == at::kMPS, "x must be on MPS device");
    TORCH_CHECK(x.dim() == 2, "x must be 2D (Batch, N)");
    
    int64_t Batch = x.size(0);
    int64_t N = x.size(1);
    
    auto out = x.clone().contiguous();
    
    if (!kernels.softmaxBatchedPSO) {
        // Fallback
        return torch::softmax(x / temperature, -1);
    }
    
    @autoreleasepool {
        auto stream = at::mps::getCurrentMPSStream();
        auto cmdBuffer = stream->commandBuffer();
        auto encoder = [cmdBuffer computeCommandEncoder];
        
        [encoder setComputePipelineState:kernels.softmaxBatchedPSO];
        [encoder setBuffer:getMTLBufferStorage(out) offset:0 atIndex:0];
        
        uint32_t N_uint = (uint32_t)N;
        uint32_t Batch_uint = (uint32_t)Batch;
        [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:1];
        [encoder setBytes:&Batch_uint length:sizeof(Batch_uint) atIndex:2];
        [encoder setBytes:&temperature length:sizeof(temperature) atIndex:3];
        
        NSUInteger shared_size = 256 * sizeof(float);
        [encoder setThreadgroupMemoryLength:shared_size atIndex:0];
        
        [encoder dispatchThreadgroups:MTLSizeMake(Batch, 1, 1) 
            threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        
        [encoder endEncoding];
        stream->synchronize(SyncType::COMMIT_AND_WAIT);
    }
    
    return out;
}

// -----------------------------------------------------------------------------
// Batched Trace
// -----------------------------------------------------------------------------

torch::Tensor trace_batched_metal(torch::Tensor A) {
    load_core_kernels();
    
    TORCH_CHECK(A.device().type() == at::kMPS, "A must be on MPS device");
    TORCH_CHECK(A.dim() == 3, "A must be 3D (Batch, N, N)");
    TORCH_CHECK(A.size(1) == A.size(2), "A must be square");
    
    int64_t Batch = A.size(0);
    int64_t N = A.size(1);
    
    auto A_contig = A.contiguous();
    auto traces = torch::zeros({Batch}, A.options());
    
    if (!kernels.traceBatchedPSO) {
        // Fallback
        return torch::sum(torch::diagonal(A, 0, -2, -1), -1);
    }
    
    @autoreleasepool {
        auto stream = at::mps::getCurrentMPSStream();
        auto cmdBuffer = stream->commandBuffer();
        auto encoder = [cmdBuffer computeCommandEncoder];
        
        [encoder setComputePipelineState:kernels.traceBatchedPSO];
        [encoder setBuffer:getMTLBufferStorage(A_contig) offset:A_contig.storage_offset() * A_contig.element_size() atIndex:0];
        [encoder setBuffer:getMTLBufferStorage(traces) offset:0 atIndex:1];
        
        uint32_t N_uint = (uint32_t)N;
        uint32_t Batch_uint = (uint32_t)Batch;
        [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:2];
        [encoder setBytes:&Batch_uint length:sizeof(Batch_uint) atIndex:3];
        
        NSUInteger shared_size = 256 * sizeof(float);
        [encoder setThreadgroupMemoryLength:shared_size atIndex:0];
        
        [encoder dispatchThreadgroups:MTLSizeMake(Batch, 1, 1) 
            threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        
        [encoder endEncoding];
        stream->synchronize(SyncType::COMMIT_AND_WAIT);
    }
    
    return traces;
}

// -----------------------------------------------------------------------------
// Column Norms
// -----------------------------------------------------------------------------

torch::Tensor column_norms_metal(torch::Tensor A) {
    load_core_kernels();
    
    TORCH_CHECK(A.device().type() == at::kMPS, "A must be on MPS device");
    TORCH_CHECK(A.dim() == 2, "A must be 2D");
    
    int64_t M = A.size(0);
    int64_t N = A.size(1);
    
    A = A.contiguous();
    auto norms = torch::empty({N}, A.options());
    
    if (!kernels.columnNormsPSO) {
        // Fallback: compute norms manually
        return (A * A).sum(0).sqrt();
    }
    
    @autoreleasepool {
        auto stream = at::mps::getCurrentMPSStream();
        auto cmdBuffer = stream->commandBuffer();
        auto encoder = [cmdBuffer computeCommandEncoder];
        
        [encoder setComputePipelineState:kernels.columnNormsPSO];
        [encoder setBuffer:getMTLBufferStorage(A) offset:A.storage_offset() * A.element_size() atIndex:0];
        [encoder setBuffer:getMTLBufferStorage(norms) offset:0 atIndex:1];
        
        uint32_t M_uint = (uint32_t)M;
        uint32_t N_uint = (uint32_t)N;
        [encoder setBytes:&M_uint length:sizeof(M_uint) atIndex:2];
        [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:3];
        
        NSUInteger tg_size = 256;
        NSUInteger shared_size = tg_size * sizeof(float);
        [encoder setThreadgroupMemoryLength:shared_size atIndex:0];
        
        // One threadgroup per column
        [encoder dispatchThreadgroups:MTLSizeMake(N, 1, 1) 
            threadsPerThreadgroup:MTLSizeMake(tg_size, 1, 1)];
        
        [encoder endEncoding];
        stream->synchronize(SyncType::COMMIT_AND_WAIT);
    }
    
    return norms;
}

// -----------------------------------------------------------------------------
// Forward Declarations for Python Bindings
// -----------------------------------------------------------------------------

torch::Tensor cholesky_batched_metal(torch::Tensor A);
torch::Tensor cholesky_solve_batched_metal(torch::Tensor L, torch::Tensor b);

// =============================================================================
// SVD KERNELS (ported from metalsvd)
// =============================================================================

struct SVDKernels {
    id<MTLFunction> jacobi = nil;
    id<MTLFunction> jacobi_icb = nil;
    id<MTLFunction> jacobi_icb_vec4 = nil;
    id<MTLFunction> jacobi_fused = nil;
    id<MTLFunction> jacobi_fused_256 = nil;
    id<MTLFunction> norm = nil;
    id<MTLFunction> normalize = nil;
    
    id<MTLComputePipelineState> jacobiPSO = nil;
    id<MTLComputePipelineState> jacobiICBPSO = nil;
    id<MTLComputePipelineState> jacobiICBVec4PSO = nil;
    id<MTLComputePipelineState> jacobiFusedPSO = nil;
    id<MTLComputePipelineState> jacobiFused256PSO = nil;
    id<MTLComputePipelineState> jacobiFused128PSO = nil;
    id<MTLComputePipelineState> jacobiFused64PSO = nil;
    id<MTLComputePipelineState> jacobiFused512PSO = nil;
    id<MTLComputePipelineState> jacobiFused1024PSO = nil;
    id<MTLComputePipelineState> normPSO = nil;
    id<MTLComputePipelineState> normalizePSO = nil;
    
    NSMutableDictionary<NSNumber*, id<MTLIndirectCommandBuffer>>* icbCache = nil;
    NSMutableDictionary<NSNumber*, id<MTLBuffer>>* stepBufferCache = nil;
};

static SVDKernels svdKFloat;
static SVDKernels svdKHalf;
static SVDKernels svdKBFloat;
static bool svd_kernels_loaded = false;

void load_svd_kernels_typed(id<MTLLibrary> lib, SVDKernels& k, NSString* suffix, bool required) {
    k.icbCache = [NSMutableDictionary new];
    k.stepBufferCache = [NSMutableDictionary new];

    k.jacobi = [lib newFunctionWithName:[NSString stringWithFormat:@"jacobi_rotate_kernel_optimized_%@", suffix]];
    k.jacobi_icb = [lib newFunctionWithName:[NSString stringWithFormat:@"jacobi_rotate_kernel_icb_%@", suffix]];
    k.jacobi_icb_vec4 = [lib newFunctionWithName:[NSString stringWithFormat:@"jacobi_rotate_kernel_vec4_clean_%@", suffix]];
    k.jacobi_fused = [lib newFunctionWithName:[NSString stringWithFormat:@"svd_fused_block_kernel_%@", suffix]];
    k.jacobi_fused_256 = [lib newFunctionWithName:[NSString stringWithFormat:@"svd_fused_block_kernel_256_%@", suffix]];
    k.norm = [lib newFunctionWithName:[NSString stringWithFormat:@"column_norm_kernel_%@", suffix]];
    k.normalize = [lib newFunctionWithName:[NSString stringWithFormat:@"normalize_kernel_%@", suffix]];
    
    NSError* error = nil;
    id<MTLDevice> device = MPSDevice::getInstance()->device();
    
    if (k.jacobi) k.jacobiPSO = [device newComputePipelineStateWithFunction:k.jacobi error:&error];
    
    if (k.jacobi_icb) {
        MTLComputePipelineDescriptor* desc = [[MTLComputePipelineDescriptor alloc] init];
        desc.computeFunction = k.jacobi_icb;
        desc.supportIndirectCommandBuffers = YES;
        k.jacobiICBPSO = [device newComputePipelineStateWithDescriptor:desc options:MTLPipelineOptionNone reflection:nil error:&error];
    }
    
    if (k.jacobi_icb_vec4) {
        MTLComputePipelineDescriptor* desc = [[MTLComputePipelineDescriptor alloc] init];
        desc.computeFunction = k.jacobi_icb_vec4;
        desc.supportIndirectCommandBuffers = YES;
        k.jacobiICBVec4PSO = [device newComputePipelineStateWithDescriptor:desc options:MTLPipelineOptionNone reflection:nil error:&error];
    }
    
    if (k.jacobi_fused) k.jacobiFusedPSO = [device newComputePipelineStateWithFunction:k.jacobi_fused error:&error];
    if (k.jacobi_fused_256) k.jacobiFused256PSO = [device newComputePipelineStateWithFunction:k.jacobi_fused_256 error:&error];
    
    id<MTLFunction> f128 = [lib newFunctionWithName:[NSString stringWithFormat:@"svd_fused_block_kernel_128_%@", suffix]];
    if (f128) k.jacobiFused128PSO = [device newComputePipelineStateWithFunction:f128 error:&error];
    
    id<MTLFunction> f64 = [lib newFunctionWithName:[NSString stringWithFormat:@"svd_fused_block_kernel_64_%@", suffix]];
    if (f64) k.jacobiFused64PSO = [device newComputePipelineStateWithFunction:f64 error:&error];
    
    id<MTLFunction> f512 = [lib newFunctionWithName:[NSString stringWithFormat:@"svd_fused_block_kernel_512_%@", suffix]];
    if (f512) k.jacobiFused512PSO = [device newComputePipelineStateWithFunction:f512 error:&error];
    
    id<MTLFunction> f1024 = [lib newFunctionWithName:[NSString stringWithFormat:@"svd_fused_block_kernel_1024_%@", suffix]];
    if (f1024) k.jacobiFused1024PSO = [device newComputePipelineStateWithFunction:f1024 error:&error];
    
    if (k.norm) k.normPSO = [device newComputePipelineStateWithFunction:k.norm error:&error];
    if (k.normalize) k.normalizePSO = [device newComputePipelineStateWithFunction:k.normalize error:&error];
}

void load_svd_kernels() {
    if (svd_kernels_loaded) return;
    load_core_kernels();  // Ensure coreLib is loaded
    
    load_svd_kernels_typed(coreLib, svdKFloat, @"float", true);
    load_svd_kernels_typed(coreLib, svdKHalf, @"half", false);
    load_svd_kernels_typed(coreLib, svdKBFloat, @"bfloat", false);
    svd_kernels_loaded = true;
}

std::pair<std::vector<int>, int> svd_generate_ordering(int N) {
    std::vector<int> all_pairs;
    int num_steps = N - 1; 
    std::vector<int> players(N);
    for(int i=0; i<N; ++i) players[i] = i;
    for(int s=0; s<num_steps; ++s) {
        for(int k=0; k<N/2; ++k) {
            all_pairs.push_back(players[k]);
            all_pairs.push_back(players[N - 1 - k]);
        }
        int last = players.back();
        for(int i=N-1; i>1; --i) players[i] = players[i-1];
        players[1] = last;
    }
    return {all_pairs, num_steps};
}

std::vector<torch::Tensor> svd_forward(torch::Tensor A) { 
    TORCH_CHECK(A.device().is_mps(), "Input tensor must be on MPS");
    load_svd_kernels();
    
    SVDKernels* kernels = nullptr;
    if (A.scalar_type() == torch::kFloat32) kernels = &svdKFloat;
    else if (A.scalar_type() == torch::kHalf) kernels = &svdKHalf;
    else if (A.scalar_type() == torch::kBFloat16) kernels = &svdKBFloat;
    else TORCH_CHECK(false, "Unsupported dtype.");
    
    if (A.dim() == 2) A = A.unsqueeze(0);
    
    int64_t Batch = A.size(0);
    int64_t M = A.size(1);
    int64_t N = A.size(2);
    
    TORCH_CHECK(N % 2 == 0, "N must be even");

    torch::Tensor V = torch::eye(N, A.options()).expand({Batch, N, N}).contiguous();
    torch::Tensor A_T = A.transpose(1, 2).contiguous(); 
    torch::Tensor V_T = V.transpose(1, 2).contiguous(); 
    
    auto [pairs_cpu, num_steps] = svd_generate_ordering(N);
    int num_pairs = N / 2;
    int threads_per_pair = 32;
    if (N >= 64) threads_per_pair = 32;
    if (N >= 128) threads_per_pair = 16;
    if (N >= 256) threads_per_pair = 8;
    if (N >= 512) threads_per_pair = 4;
    if (N >= 1024) threads_per_pair = 2;
    
    int specialized_mode = 0;
    if (N == 1024) specialized_mode = 5;
    else if (N == 512) specialized_mode = 4;
    else if (N == 256) specialized_mode = 1;
    else if (N == 128) specialized_mode = 2;
    else if (N == 64) specialized_mode = 3;
    
    bool use_fused_any = (specialized_mode > 0) || (N <= 256);

    torch::Tensor PairsTens = torch::tensor(pairs_cpu, torch::dtype(torch::kInt32).device(torch::kCPU)).contiguous();
    PairsTens = PairsTens.to(A.device());
    
    MPSStream* stream = getCurrentMPSStream();
    id<MTLComputeCommandEncoder> encoder = stream->commandEncoder();
    
    uint32_t BatchStrideA = (uint32_t)(N * M);
    uint32_t BatchStrideV = (uint32_t)(N * N);
    uint32_t M_u = (uint32_t)M;
    uint32_t N_u = (uint32_t)N;

    if (use_fused_any) {
        id<MTLComputePipelineState> fusedPSO = kernels->jacobiFusedPSO;
        if (specialized_mode == 5) fusedPSO = kernels->jacobiFused1024PSO;
        else if (specialized_mode == 4) fusedPSO = kernels->jacobiFused512PSO;
        else if (specialized_mode == 1) fusedPSO = kernels->jacobiFused256PSO;
        else if (specialized_mode == 2) fusedPSO = kernels->jacobiFused128PSO;
        else if (specialized_mode == 3) fusedPSO = kernels->jacobiFused64PSO;
        
        TORCH_CHECK(fusedPSO, "Failed to get Fused PSO");
        
        [encoder setComputePipelineState:fusedPSO];
        mtl_setBuffer(encoder, A_T, 0);
        mtl_setBuffer(encoder, V_T, 1);
        [encoder setBuffer:getMTLBufferStorage(PairsTens) offset:(PairsTens.storage_offset() * PairsTens.element_size()) atIndex:2];
        
        if (specialized_mode > 0) {
            [encoder setBytes:&M_u length:sizeof(uint32_t) atIndex:3];
            [encoder setBytes:&BatchStrideA length:sizeof(uint32_t) atIndex:4];
            [encoder setBytes:&BatchStrideV length:sizeof(uint32_t) atIndex:5];
        } else {
            [encoder setBytes:&M_u length:sizeof(uint32_t) atIndex:3];
            [encoder setBytes:&N_u length:sizeof(uint32_t) atIndex:4];
            uint32_t NumPairs_u = (uint32_t)num_pairs;
            [encoder setBytes:&NumPairs_u length:sizeof(uint32_t) atIndex:5];
            uint32_t NumSteps_u = (uint32_t)num_steps;
            [encoder setBytes:&NumSteps_u length:sizeof(uint32_t) atIndex:6];
            uint32_t TPP_u = (uint32_t)threads_per_pair;
            [encoder setBytes:&TPP_u length:sizeof(uint32_t) atIndex:7];
            [encoder setBytes:&BatchStrideA length:sizeof(uint32_t) atIndex:8];
            [encoder setBytes:&BatchStrideV length:sizeof(uint32_t) atIndex:9];
        }
        
        int total_threads = num_pairs * threads_per_pair;
        [encoder dispatchThreadgroups:MTLSizeMake(1, 1, Batch) threadsPerThreadgroup:MTLSizeMake(total_threads, 1, 1)];
    } else {
        // ICB path for large matrices - simplified version
        int sweeps = 6;
        id<MTLComputePipelineState> rotatePSO = kernels->jacobiICBPSO ? kernels->jacobiICBPSO : kernels->jacobiPSO;
        TORCH_CHECK(rotatePSO, "No rotate PSO available");
        
        int threads_per_group = std::min((int)rotatePSO.maxTotalThreadsPerThreadgroup, 256);
        int elem_size = A.element_size();
        NSUInteger sharedMemSize = ((threads_per_group + 31) / 32) * 3 * elem_size;
        
        [encoder setComputePipelineState:rotatePSO];
        mtl_setBuffer(encoder, A_T, 0);
        mtl_setBuffer(encoder, V_T, 1);
        [encoder setBytes:&M_u length:sizeof(uint32_t) atIndex:3];
        [encoder setBytes:&N_u length:sizeof(uint32_t) atIndex:4];
        [encoder setBytes:&BatchStrideA length:sizeof(uint32_t) atIndex:5];
        [encoder setBytes:&BatchStrideV length:sizeof(uint32_t) atIndex:6];
        [encoder setThreadgroupMemoryLength:sharedMemSize atIndex:0];
        
        for (int sw = 0; sw < sweeps; ++sw) {
            for (int step = 0; step < num_steps; ++step) {
                size_t pairs_offset = step * num_pairs * sizeof(int) * 2;
                [encoder setBuffer:getMTLBufferStorage(PairsTens) 
                            offset:(PairsTens.storage_offset() * PairsTens.element_size() + pairs_offset) 
                           atIndex:2];
                [encoder dispatchThreadgroups:MTLSizeMake(num_pairs, 1, Batch) 
                    threadsPerThreadgroup:MTLSizeMake(threads_per_group, 1, 1)];
            }
        }
    }
    
    torch::Tensor S = torch::empty({Batch, N}, A.options()); 
    
    [encoder setComputePipelineState:kernels->normPSO];
    mtl_setBuffer(encoder, A_T, 0);
    mtl_setBuffer(encoder, S, 1);
    [encoder setBytes:&M_u length:sizeof(uint32_t) atIndex:2];
    [encoder setBytes:&N_u length:sizeof(uint32_t) atIndex:3];
    [encoder setBytes:&BatchStrideA length:sizeof(uint32_t) atIndex:4];
    uint32_t BatchStrideS = (uint32_t)N;
    [encoder setBytes:&BatchStrideS length:sizeof(uint32_t) atIndex:5];
    
    MTLSize normGridSize = MTLSizeMake(N, Batch, 1);
    MTLSize normGroupSize = MTLSizeMake(std::min((int)N, (int)kernels->normPSO.maxTotalThreadsPerThreadgroup), 1, 1);
    [encoder dispatchThreads:normGridSize threadsPerThreadgroup:normGroupSize];

    torch::Tensor U_T = torch::empty_like(A_T);
    
    [encoder setComputePipelineState:kernels->normalizePSO];
    mtl_setBuffer(encoder, A_T, 0);
    mtl_setBuffer(encoder, S, 1);
    mtl_setBuffer(encoder, U_T, 2);
    [encoder setBytes:&M_u length:sizeof(uint32_t) atIndex:3];
    [encoder setBytes:&N_u length:sizeof(uint32_t) atIndex:4];
    [encoder setBytes:&BatchStrideA length:sizeof(uint32_t) atIndex:5];
    [encoder setBytes:&BatchStrideS length:sizeof(uint32_t) atIndex:6];
    
    [encoder dispatchThreads:normGridSize threadsPerThreadgroup:normGroupSize];
    
    return {U_T.transpose(1, 2).contiguous(), S, V_T.transpose(1, 2).contiguous()};
}

// =============================================================================
// EIGH KERNELS (ported from metaleig)
// =============================================================================

struct EighKernels {
    id<MTLFunction> jacobi = nil;
    id<MTLFunction> dot_columns = nil;
    id<MTLFunction> fused_generic = nil;
    id<MTLFunction> fused_64 = nil;
    id<MTLFunction> fused_128 = nil;
    id<MTLFunction> fused_256 = nil;
    id<MTLFunction> fused_512 = nil;
    id<MTLFunction> fused_1024 = nil;
    
    id<MTLComputePipelineState> jacobiPSO = nil;
    id<MTLComputePipelineState> dotColumnsPSO = nil;
    id<MTLComputePipelineState> fusedGenericPSO = nil;
    id<MTLComputePipelineState> fused64PSO = nil;
    id<MTLComputePipelineState> fused128PSO = nil;
    id<MTLComputePipelineState> fused256PSO = nil;
    id<MTLComputePipelineState> fused512PSO = nil;
    id<MTLComputePipelineState> fused1024PSO = nil;
    
    id<MTLFunction> jacobiICB = nil;
    id<MTLComputePipelineState> jacobiICBPSO = nil;
    NSMutableDictionary<NSNumber*, id<MTLIndirectCommandBuffer>>* icbCache = nil;
    NSMutableDictionary<NSNumber*, id<MTLBuffer>>* stepBufferCache = nil;
    NSMutableDictionary<NSNumber*, id<MTLBuffer>>* uniformBufferCache = nil;
};

static EighKernels eighKFloat;
static EighKernels eighKHalf;
static EighKernels eighKBFloat;
static bool eigh_kernels_loaded = false;

void load_eigh_kernels_typed(id<MTLLibrary> lib, EighKernels& k, NSString* suffix, bool required) {
    k.icbCache = [NSMutableDictionary new];
    k.stepBufferCache = [NSMutableDictionary new];
    k.uniformBufferCache = [NSMutableDictionary new];
    
    NSError* error = nil;
    id<MTLDevice> device = MPSDevice::getInstance()->device();

    k.jacobi = [lib newFunctionWithName:[NSString stringWithFormat:@"jacobi_rotate_kernel_optimized_%@", suffix]];
    k.dot_columns = [lib newFunctionWithName:[NSString stringWithFormat:@"dot_columns_kernel_%@", suffix]];
    k.fused_generic = [lib newFunctionWithName:[NSString stringWithFormat:@"svd_fused_block_kernel_generic_%@", suffix]];
    k.fused_64 = [lib newFunctionWithName:[NSString stringWithFormat:@"svd_fused_block_kernel_64_%@", suffix]];
    k.fused_128 = [lib newFunctionWithName:[NSString stringWithFormat:@"svd_fused_block_kernel_128_%@", suffix]];
    k.fused_256 = [lib newFunctionWithName:[NSString stringWithFormat:@"svd_fused_block_kernel_256_%@", suffix]];
    k.fused_512 = [lib newFunctionWithName:[NSString stringWithFormat:@"svd_fused_block_kernel_512_%@", suffix]];
    k.fused_1024 = [lib newFunctionWithName:[NSString stringWithFormat:@"svd_fused_block_kernel_1024_%@", suffix]];
    
    k.jacobiICB = [lib newFunctionWithName:[NSString stringWithFormat:@"jacobi_rotate_kernel_icb_%@", suffix]];
    if (k.jacobiICB) {
        MTLComputePipelineDescriptor* desc = [[MTLComputePipelineDescriptor alloc] init];
        desc.computeFunction = k.jacobiICB;
        desc.supportIndirectCommandBuffers = YES;
        k.jacobiICBPSO = [device newComputePipelineStateWithDescriptor:desc options:MTLPipelineOptionNone reflection:nil error:&error];
    }
    
    if (k.jacobi) k.jacobiPSO = [device newComputePipelineStateWithFunction:k.jacobi error:&error];
    if (k.dot_columns) k.dotColumnsPSO = [device newComputePipelineStateWithFunction:k.dot_columns error:&error];
    if (k.fused_generic) k.fusedGenericPSO = [device newComputePipelineStateWithFunction:k.fused_generic error:&error];
    if (k.fused_64) k.fused64PSO = [device newComputePipelineStateWithFunction:k.fused_64 error:&error];
    if (k.fused_128) k.fused128PSO = [device newComputePipelineStateWithFunction:k.fused_128 error:&error];
    if (k.fused_256) k.fused256PSO = [device newComputePipelineStateWithFunction:k.fused_256 error:&error];
    if (k.fused_512) k.fused512PSO = [device newComputePipelineStateWithFunction:k.fused_512 error:&error];
    if (k.fused_1024) k.fused1024PSO = [device newComputePipelineStateWithFunction:k.fused_1024 error:&error];
}

void load_eigh_kernels() {
    if (eigh_kernels_loaded) return;
    load_core_kernels();
    
    load_eigh_kernels_typed(coreLib, eighKFloat, @"float", true);
    load_eigh_kernels_typed(coreLib, eighKHalf, @"half", false);
    load_eigh_kernels_typed(coreLib, eighKBFloat, @"bfloat", false);
    eigh_kernels_loaded = true;
}

std::vector<torch::Tensor> eigh_forward(torch::Tensor A) { 
    TORCH_CHECK(A.device().is_mps(), "Input tensor must be on MPS");
    load_eigh_kernels();
    
    EighKernels* kernels = nullptr;
    if (A.scalar_type() == torch::kFloat32) kernels = &eighKFloat;
    else if (A.scalar_type() == torch::kHalf) kernels = &eighKHalf;
    else if (A.scalar_type() == torch::kBFloat16) kernels = &eighKBFloat;
    else TORCH_CHECK(false, "Unsupported dtype.");
    
    if (A.dim() == 2) A = A.unsqueeze(0);
    
    int64_t Batch = A.size(0);
    int64_t M = A.size(1);
    int64_t N = A.size(2);
    
    TORCH_CHECK(N % 2 == 0, "N must be even");

    torch::Tensor V = torch::eye(N, A.options()).expand({Batch, N, N}).contiguous();
    torch::Tensor A_T = A.transpose(1, 2).contiguous(); 
    torch::Tensor V_T = V.transpose(1, 2).contiguous(); 
    
    auto [pairs_cpu, num_steps] = svd_generate_ordering(N);
    int num_pairs = N / 2;
    int threads_per_pair = std::min(32, std::max(1, 1024 / num_pairs));
    
    torch::Tensor PairsTens = torch::tensor(pairs_cpu, torch::dtype(torch::kInt32).device(torch::kCPU)).contiguous();
    PairsTens = PairsTens.to(A.device());
    
    MPSStream* stream = getCurrentMPSStream();
    id<MTLComputeCommandEncoder> encoder = stream->commandEncoder();
    
    uint32_t BatchStrideA = (uint32_t)(N * M);
    uint32_t BatchStrideV = (uint32_t)(N * N);
    uint32_t M_u = (uint32_t)M;
    uint32_t N_u = (uint32_t)N;
    
    id<MTLComputePipelineState> fusedPSO = nil;
    if (N == 64 && kernels->fused64PSO) fusedPSO = kernels->fused64PSO;
    else if (N == 128 && kernels->fused128PSO) fusedPSO = kernels->fused128PSO;
    
    if (fusedPSO) {
        [encoder setComputePipelineState:fusedPSO];
        mtl_setBuffer(encoder, A_T, 0);
        mtl_setBuffer(encoder, V_T, 1);
        [encoder setBuffer:getMTLBufferStorage(PairsTens) offset:(PairsTens.storage_offset() * PairsTens.element_size()) atIndex:2];
        [encoder setBytes:&M_u length:sizeof(uint32_t) atIndex:3];
        [encoder setBytes:&BatchStrideA length:sizeof(uint32_t) atIndex:4];
        [encoder setBytes:&BatchStrideV length:sizeof(uint32_t) atIndex:5];
        [encoder dispatchThreadgroups:MTLSizeMake(1, 1, Batch) threadsPerThreadgroup:MTLSizeMake(1024, 1, 1)];
    } else {
        // Iterative fallback
        int sweeps = 15;
        id<MTLComputePipelineState> rotatePSO = kernels->jacobiPSO;
        TORCH_CHECK(rotatePSO, "No jacobi PSO");
        
        int threads_per_group = std::min((int)rotatePSO.maxTotalThreadsPerThreadgroup, 256);
        int elem_size = A.element_size();
        NSUInteger sharedMemSize = ((threads_per_group + 31) / 32) * 3 * elem_size;
        
        [encoder setComputePipelineState:rotatePSO];
        mtl_setBuffer(encoder, A_T, 0);
        mtl_setBuffer(encoder, V_T, 1);
        [encoder setBytes:&M_u length:sizeof(uint32_t) atIndex:3];
        [encoder setBytes:&N_u length:sizeof(uint32_t) atIndex:4];
        [encoder setBytes:&BatchStrideA length:sizeof(uint32_t) atIndex:5];
        [encoder setBytes:&BatchStrideV length:sizeof(uint32_t) atIndex:6];
        [encoder setThreadgroupMemoryLength:sharedMemSize atIndex:0];
        
        for (int sw = 0; sw < sweeps; ++sw) {
            for (int step = 0; step < num_steps; ++step) {
                size_t pairs_offset = step * num_pairs * sizeof(int) * 2;
                [encoder setBuffer:getMTLBufferStorage(PairsTens) 
                            offset:(PairsTens.storage_offset() * PairsTens.element_size() + pairs_offset) 
                           atIndex:2];
                [encoder dispatchThreadgroups:MTLSizeMake(num_pairs, 1, Batch) 
                    threadsPerThreadgroup:MTLSizeMake(threads_per_group, 1, 1)];
            }
        }
    }
    
    torch::Tensor Eigenvalues = torch::empty({Batch, N}, A.options());
    
    id<MTLComputePipelineState> dotPSO = kernels->dotColumnsPSO;
    TORCH_CHECK(dotPSO, "No dot columns PSO");
    
    [encoder setComputePipelineState:dotPSO];
    mtl_setBuffer(encoder, A_T, 0);
    mtl_setBuffer(encoder, V_T, 1);
    mtl_setBuffer(encoder, Eigenvalues, 2);
    [encoder setBytes:&M_u length:sizeof(uint32_t) atIndex:3];
    [encoder setBytes:&N_u length:sizeof(uint32_t) atIndex:4];
    [encoder setBytes:&BatchStrideA length:sizeof(uint32_t) atIndex:5];
    [encoder setBytes:&BatchStrideV length:sizeof(uint32_t) atIndex:6];
    uint32_t BatchStrideE = (uint32_t)N;
    [encoder setBytes:&BatchStrideE length:sizeof(uint32_t) atIndex:7];
    
    int dot_tpg = std::min((int)dotPSO.maxTotalThreadsPerThreadgroup, 256);
    NSUInteger dotSharedMem = ((dot_tpg + 31) / 32) * sizeof(float);
    [encoder setThreadgroupMemoryLength:dotSharedMem atIndex:0];
    
    [encoder dispatchThreadgroups:MTLSizeMake(N, 1, Batch) threadsPerThreadgroup:MTLSizeMake(dot_tpg, 1, 1)];
    
    return {Eigenvalues, V_T.transpose(1, 2).contiguous()};
}

// -----------------------------------------------------------------------------
// Python Bindings
// -----------------------------------------------------------------------------

PYBIND11_MODULE(metalcore_backend, m) {
    m.def("trsm", &trsm_metal, "Triangular solve on Metal",
          py::arg("A"), py::arg("b"), py::arg("lower") = true, py::arg("transpose") = false);
    
    m.def("geqr2", &geqr2_metal, "Panel QR factorization (unblocked Householder)",
          py::arg("A"));
    
    m.def("larfb", &larfb_metal, "Apply block Householder reflector",
          py::arg("C"), py::arg("V"), py::arg("T"), py::arg("trans") = true, py::arg("panel_start") = 0);
    
    m.def("larft", &larft_metal, "Build triangular factor T for WY representation",
          py::arg("V"), py::arg("tau"), py::arg("panel_start") = 0);
    
    m.def("qr", &qr_metal, "Blocked Householder QR decomposition",
          py::arg("A"), py::arg("block_size") = 32);
    
    m.def("qr_fused", &qr_fused_metal, "Fully fused QR (single Metal dispatch)",
          py::arg("A"));
    
    m.def("qr_batched", &qr_batched_metal, "Batched QR (parallel matrices in single dispatch)",
          py::arg("A_batch"));
    
    m.def("trsm_batched", &trsm_batched_metal, "Batched triangular solve",
          py::arg("R"), py::arg("B"));
    
    m.def("solve_batched", &solve_batched_metal, "Fused batched solve (QR + Q.T@b + TRSM)",
          py::arg("A"), py::arg("b"));
    
    m.def("column_norms", &column_norms_metal, "Compute column norms",
          py::arg("A"));
    
    m.def("cholesky_batched", &cholesky_batched_metal, "Batched Cholesky decomposition",
          py::arg("A"));
    
    m.def("cholesky_solve_batched", &cholesky_solve_batched_metal, "Batched Cholesky solve",
          py::arg("L"), py::arg("b"));
    
    // SVD (ported from metalsvd)
    m.def("svd_forward", &svd_forward, "Batched SVD (Metal Jacobi)",
          py::arg("A"));
    
    // Eigendecomposition (ported from metaleig)
    m.def("eigh_forward", &eigh_forward, "Batched Eigendecomposition (Metal Jacobi)",
          py::arg("A"));
    
    // New optimization kernels
    m.def("column_norm_sort", &column_norm_sort_metal, "De Rijk column sorting for SVD optimization",
          py::arg("A"));
    
    m.def("sign_canonicalize", &sign_canonicalize_metal, "SVD sign canonicalization (in-place)",
          py::arg("U"), py::arg("V"));
    
    m.def("batched_qt_b", &batched_qt_b_metal, "Batched Q.T @ b for fused solve",
          py::arg("Q"), py::arg("b"));
    
    // High-impact ML/LA kernels
    m.def("lu_batched", &lu_batched_metal, "Batched LU decomposition with pivoting",
          py::arg("A"));
    
    m.def("syrk_batched", &syrk_batched_metal, "Batched SYRK: C = A.T @ A",
          py::arg("A"));
    
    m.def("frobenius_norm_batched", &frobenius_norm_batched_metal, "Batched Frobenius norm",
          py::arg("A"));
    
    m.def("softmax_batched", &softmax_batched_metal, "Batched softmax with temperature",
          py::arg("x"), py::arg("temperature") = 1.0f);
    
    m.def("trace_batched", &trace_batched_metal, "Batched trace (sum of diagonal)",
          py::arg("A"));
}

// -----------------------------------------------------------------------------
// Batched Cholesky Decomposition
// -----------------------------------------------------------------------------

torch::Tensor cholesky_batched_metal(torch::Tensor A) {
    load_core_kernels();
    
    TORCH_CHECK(A.device().type() == at::kMPS, "A must be on MPS device");
    TORCH_CHECK(A.dim() == 3, "A must be 3D (batch, N, N)");
    TORCH_CHECK(A.size(1) == A.size(2), "A must be square");
    
    int64_t batch_size = A.size(0);
    int64_t N = A.size(1);
    
    // Clone and make contiguous (we modify in-place)
    auto L = A.clone().contiguous();
    
    if (!kernels.choleskyBatchedPSO) {
        // CPU fallback using manual Cholesky
        auto L_cpu = L.cpu();
        for (int64_t i = 0; i < batch_size; i++) {
            L_cpu[i] = at::linalg_cholesky(L_cpu[i]);
        }
        return L_cpu.to(A.device());
    }
    
    @autoreleasepool {
        auto stream = at::mps::getCurrentMPSStream();
        auto cmdBuffer = stream->commandBuffer();
        auto encoder = [cmdBuffer computeCommandEncoder];
        
        [encoder setComputePipelineState:kernels.choleskyBatchedPSO];
        [encoder setBuffer:getMTLBufferStorage(L) offset:L.storage_offset() * L.element_size() atIndex:0];
        
        uint32_t N_uint = (uint32_t)N;
        uint32_t batch_uint = (uint32_t)batch_size;
        [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:1];
        [encoder setBytes:&batch_uint length:sizeof(batch_uint) atIndex:2];
        
        // Shared memory for N*N panel (MAGMA-style optimization)
        NSUInteger shared_size = N * N * sizeof(float);
        [encoder setThreadgroupMemoryLength:shared_size atIndex:0];
        
        // One threadgroup per batch, up to 64 threads per group
        NSUInteger tg_size = std::min((NSUInteger)64, std::max((NSUInteger)32, (NSUInteger)N));
        [encoder dispatchThreadgroups:MTLSizeMake(batch_size, 1, 1) 
            threadsPerThreadgroup:MTLSizeMake(tg_size, 1, 1)];
        
        [encoder endEncoding];
        stream->synchronize(SyncType::COMMIT_AND_WAIT);
    }
    
    return L;
}

// -----------------------------------------------------------------------------
// Batched Cholesky Solve (L @ L.T @ x = b)
// -----------------------------------------------------------------------------

torch::Tensor cholesky_solve_batched_metal(torch::Tensor L, torch::Tensor b) {
    load_core_kernels();
    
    TORCH_CHECK(L.device().type() == at::kMPS, "L must be on MPS device");
    TORCH_CHECK(b.device().type() == at::kMPS, "b must be on MPS device");
    TORCH_CHECK(L.dim() == 3, "L must be 3D (batch, N, N)");
    
    int64_t batch_size = L.size(0);
    int64_t N = L.size(1);
    int64_t K = b.dim() == 3 ? b.size(2) : 1;
    
    L = L.contiguous();
    auto x = b.clone().contiguous();
    if (x.dim() == 2) {
        x = x.unsqueeze(-1);  // (batch, N) -> (batch, N, 1)
    }
    
    if (!kernels.choleskySolveBatchedPSO) {
        // CPU fallback
        auto L_cpu = L.cpu();
        auto x_cpu = x.cpu();
        for (int64_t i = 0; i < batch_size; i++) {
            x_cpu[i] = at::cholesky_solve(x_cpu[i], L_cpu[i]);
        }
        return b.dim() == 2 ? x_cpu.squeeze(-1).to(L.device()) : x_cpu.to(L.device());
    }
    
    // Single fused kernel: forward + back substitution with zero-copy transpose
    // The kernel accesses L[j,i] directly for L.T[i,j] - no memory copy needed
    @autoreleasepool {
        auto stream = at::mps::getCurrentMPSStream();
        auto cmdBuffer = stream->commandBuffer();
        auto encoder = [cmdBuffer computeCommandEncoder];
        
        [encoder setComputePipelineState:kernels.choleskySolveBatchedPSO];
        [encoder setBuffer:getMTLBufferStorage(L) offset:L.storage_offset() * L.element_size() atIndex:0];
        [encoder setBuffer:getMTLBufferStorage(x) offset:x.storage_offset() * x.element_size() atIndex:1];
        
        uint32_t N_uint = (uint32_t)N;
        uint32_t K_uint = (uint32_t)K;
        uint32_t batch_uint = (uint32_t)batch_size;
        [encoder setBytes:&N_uint length:sizeof(N_uint) atIndex:2];
        [encoder setBytes:&K_uint length:sizeof(K_uint) atIndex:3];
        [encoder setBytes:&batch_uint length:sizeof(batch_uint) atIndex:4];
        
        // Shared memory for row cache (N floats)
        [encoder setThreadgroupMemoryLength:N * sizeof(float) atIndex:0];
        
        NSUInteger tg_size = std::min((NSUInteger)64, std::max((NSUInteger)K, (NSUInteger)32));
        [encoder dispatchThreadgroups:MTLSizeMake(batch_size, 1, 1) 
            threadsPerThreadgroup:MTLSizeMake(tg_size, 1, 1)];
        
        [encoder endEncoding];
        stream->synchronize(SyncType::COMMIT_AND_WAIT);
    }
    
    return b.dim() == 2 ? x.squeeze(-1) : x;
}
