// Scalar correctness primitives for the native numerical contract.
// Scalar SFU instructions and one explicit BF16 MMA accumulation primitive.
// No FP4/FP8, TMA, CUTLASS or TileLang dependency is used.
#include <torch/extension.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>

__global__ void unary_kernel(const float* x,float* y,int64_t n,int op) {
    int64_t i=int64_t(blockIdx.x)*blockDim.x+threadIdx.x;
    if(i>=n) return;
    float a=x[i],b=0;
    switch(op) {
        case 0: asm("tanh.approx.f32 %0, %1;":"=f"(b):"f"(a)); break;
        case 1: b=tanhf(a); break;
        case 2: b=cosf(a); break;
        case 3: b=sinf(a); break;
        case 4: b=expf(a); break;
        case 5: b=1.f/(1.f+expf(-a)); break;
    }
    y[i]=b;
}

torch::Tensor unary(torch::Tensor x,int64_t op) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type()==torch::kFloat32 && x.is_contiguous(),"Contiguous CUDA FP32 required");
    TORCH_CHECK(op>=0 && op<=5,"Unknown math primitive");
    c10::cuda::CUDAGuard guard(x.device());
    auto y=torch::empty_like(x);
    if(x.numel()) unary_kernel<<<(x.numel()+255)/256,256,0,c10::cuda::getCurrentCUDAStream().stream()>>>(
        x.data_ptr<float>(),y.data_ptr<float>(),x.numel(),op);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}

__device__ unsigned pair_bits(const __nv_bfloat16* x,int64_t a,int64_t b) {
    return unsigned(__bfloat16_as_ushort(x[a])) | (unsigned(__bfloat16_as_ushort(x[b]))<<16);
}

__global__ void mma_kernel(const __nv_bfloat16* a,const __nv_bfloat16* b,
                          const float* initial,float* output,int m,int n,int k) {
    int lane=threadIdx.x,group=lane>>2,tid=lane&3;
    int row=blockIdx.x*16+group,col=blockIdx.y*8+tid*2;
    int64_t batch=blockIdx.z;
    a+=batch*m*k;b+=batch*k*n;output+=batch*m*n;
    float c0=0,c1=0,c2=0,c3=0;
    if(initial) {
        initial+=batch*m*n;
        c0=initial[row*n+col];c1=initial[row*n+col+1];
        c2=initial[(row+8)*n+col];c3=initial[(row+8)*n+col+1];
    }
    // Keep exactly the ascending 16-element K tiles of the native scan.
    for(int ki=0;ki<k;ki+=16) {
        int ka=ki+tid*2;
        unsigned a0=pair_bits(a,row*k+ka,row*k+ka+1);
        unsigned a1=pair_bits(a,(row+8)*k+ka,(row+8)*k+ka+1);
        unsigned a2=pair_bits(a,row*k+ka+8,row*k+ka+9);
        unsigned a3=pair_bits(a,(row+8)*k+ka+8,(row+8)*k+ka+9);
        int cb=blockIdx.y*8+group;
        unsigned b0=pair_bits(b,ka*n+cb,(ka+1)*n+cb);
        unsigned b1=pair_bits(b,(ka+8)*n+cb,(ka+9)*n+cb);
        asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
            "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"
            : "+f"(c0),"+f"(c1),"+f"(c2),"+f"(c3)
            : "r"(a0),"r"(a1),"r"(a2),"r"(a3),"r"(b0),"r"(b1));
    }
    output[row*n+col]=c0;output[row*n+col+1]=c1;
    output[(row+8)*n+col]=c2;output[(row+8)*n+col+1]=c3;
}

// Canonical B200 BF16 block-FMA arithmetic for pre-Blackwell devices.
// Each 16-product instruction aligns the products and C to a common exponent,
// uses unnormalized operand-product exponents and 26 alignment bits, sums the
// aligned signed integers, then rounds the normalized result toward zero.
// BF16 products are exact in binary64 even outside the FP32 exponent range.
__global__ void mma_canonical_kernel(const __nv_bfloat16* a,const __nv_bfloat16* b,
                                    const float* initial,float* output,int m,int n,int k,int batches) {
    int64_t index=int64_t(blockIdx.x)*blockDim.x+threadIdx.x;
    int64_t count=int64_t(batches)*m*n;
    if(index>=count)return;
    int col=index%n,row=(index/n)%m,batch=index/(int64_t(m)*n);
    a+=int64_t(batch)*m*k;b+=int64_t(batch)*k*n;
    float acc=initial?initial[index]:0.f;
    for(int start=0;start<k;start+=16) {
        double terms[17];terms[0]=double(acc);
        int exponent=acc==0.f?-10000:int((__double_as_longlong(fabs(terms[0]))>>52)&2047)-1022;
        #pragma unroll
        for(int j=0;j<16;j++) {
            double av=double(__bfloat162float(a[row*k+start+j]));
            double bv=double(__bfloat162float(b[(start+j)*n+col]));
            terms[j+1]=av*bv;
            if(av!=0. && bv!=0.) {
                int ae=int((__double_as_longlong(fabs(av))>>52)&2047)-1022;
                int be=int((__double_as_longlong(fabs(bv))>>52)&2047)-1022;
                exponent=max(exponent,ae+be-1);
            }
        }
        if(exponent==-10000) {acc=0.f;continue;}
        long long total=0;
        #pragma unroll
        for(int j=0;j<17;j++)total+=__double2ll_rz(ldexp(terms[j],26-exponent));
        acc=__double2float_rz(ldexp(double(total),exponent-26));
    }
    output[index]=acc;
}

__device__ long long aligned_integer(unsigned long long magnitude,int shift,bool negative) {
    if(!magnitude)return 0;
    unsigned long long value=shift>=0?(magnitude<<shift):(-shift>=64?0:(magnitude>>(-shift)));
    return negative?-static_cast<long long>(value):static_cast<long long>(value);
}
__device__ float fixed_to_fp32_rz(long long total,int unit_exponent) {
    if(total==0)return 0.f;
    bool negative=total<0;unsigned long long magnitude=negative?-total:total;
    int top=63-__clzll(magnitude),exponent=top+unit_exponent;
    unsigned bits;
    if(exponent>127)bits=0x7f7fffff;
    else if(exponent>=-126) {
        int shift=top-23;
        unsigned mantissa=shift>=0?unsigned(magnitude>>shift):unsigned(magnitude<<(-shift));
        bits=unsigned(exponent+127)<<23 | (mantissa&0x7fffff);
    } else {
        int shift=unit_exponent+149;
        bits=shift>=0?unsigned(magnitude<<shift):(-shift>=64?0:unsigned(magnitude>>(-shift)));
    }
    return __uint_as_float(bits|(negative?0x80000000u:0));
}
__device__ int bf16_frexp_exponent(unsigned word) {
    unsigned exponent=(word>>7)&255,mantissa=word&127;
    return exponent?int(exponent)-126:(mantissa?32-__clz(mantissa)-133:-10000);
}
__global__ void mma_integer_kernel(const unsigned short* a,const unsigned short* b,const float* initial,
                                  float* output,int m,int n,int k,int batches) {
    int64_t index=int64_t(blockIdx.x)*blockDim.x+threadIdx.x;
    if(index>=int64_t(batches)*m*n)return;
    int col=index%n,row=(index/n)%m,batch=index/(int64_t(m)*n);
    a+=int64_t(batch)*m*k;b+=int64_t(batch)*k*n;
    float acc=initial?initial[index]:0.f;
    for(int start=0;start<k;start+=16) {
        unsigned aw[16],bw[16];unsigned c=__float_as_uint(acc),ce=(c>>23)&255,cm=c&0x7fffff;
        int exponent=ce?int(ce)-126:(cm?32-__clz(cm)-149:-10000);
        #pragma unroll
        for(int j=0;j<16;j++) {
            aw[j]=a[row*k+start+j];bw[j]=b[(start+j)*n+col];
            if((aw[j]&0x7fff) && (bw[j]&0x7fff))
                exponent=max(exponent,bf16_frexp_exponent(aw[j])+bf16_frexp_exponent(bw[j])-1);
        }
        if(exponent==-10000){acc=0.f;continue;}
        int unit=exponent-26;
        long long total=aligned_integer(ce?cm|0x800000:cm,(ce?int(ce)-150:-149)-unit,c>>31);
        #pragma unroll
        for(int j=0;j<16;j++) {
            unsigned ae=(aw[j]>>7)&255,be=(bw[j]>>7)&255;
            unsigned am=(aw[j]&127)|(ae?128:0),bm=(bw[j]&127)|(be?128:0);
            int shift=(ae?int(ae)-134:-133)+(be?int(be)-134:-133)-unit;
            total+=aligned_integer(am*bm,shift,((aw[j]^bw[j])>>15)!=0);
        }
        acc=fixed_to_fp32_rz(total,unit);
    }
    output[index]=acc;
}
__global__ void dense_fma_kernel(const float* a,const float* w,float* output,int m,int n,int k,int partitions) {
    int64_t index=int64_t(blockIdx.x)*blockDim.x+threadIdx.x;
    if(index>=int64_t(m)*n)return;
    int row=index/n,col=index%n;float acc=0.f;int size=k/partitions;
    for(int part=0;part<partitions;part++) {
        float partial=0.f;
        for(int j=part*size;j<(part+1)*size;j++)partial=__fmaf_rn(a[row*k+j],w[col*k+j],partial);
        acc=__fadd_rn(acc,partial);
    }
    output[index]=acc;
}
// One warp per output; explicit register groups and shuffle reduction order.
__global__ void bmm_reference_kernel(const float* a,const float* b,float* output,int m,int n,int k,int batches,int streams,int stride) {
    int lane=threadIdx.x&31;int64_t index=int64_t(blockIdx.x)*4+(threadIdx.x>>5);
    if(index>=int64_t(batches)*m*n)return;
    int batch=index/(int64_t(m)*n),row=(index/n)%m,col=index%n;
    a+=int64_t(batch)*m*k;b+=int64_t(batch)*k*n;
    float acc=0.f;
    if(lane<streams) {
        for(int start=lane*stride;start<k;start+=streams*stride)
            for(int j=start;j<min(k,start+stride);j++)acc=__fmaf_rn(a[row*k+j],b[j*n+col],acc);
    }
    for(int offset=streams/2;offset;offset>>=1) {
        float other=__shfl_down_sync(0xffffffff,acc,offset);
        if(lane<offset)acc=__fadd_rn(acc,other);
    }
    if(lane==0)output[index]=acc;
}
torch::Tensor bmm_reference(torch::Tensor a,torch::Tensor b,int64_t streams,int64_t stride) {
    TORCH_CHECK(a.is_cuda() && b.device()==a.device() && a.dim()==3 && b.dim()==3 && a.size(0)==b.size(0) && a.size(2)==b.size(1),"Matching CUDA batched matrices required");
    TORCH_CHECK(a.scalar_type()==torch::kFloat32 && b.scalar_type()==torch::kFloat32 && a.is_contiguous() && b.is_contiguous(),"Contiguous FP32 batched matrices required");
    TORCH_CHECK(streams>0 && streams<=32 && !(streams&(streams-1)) && stride>0,"Power-of-two warp register groups required");
    c10::cuda::CUDAGuard guard(a.device());auto out=torch::empty({a.size(0),a.size(1),b.size(2)},a.options());
    if(out.numel())bmm_reference_kernel<<<(out.numel()+3)/4,128,0,c10::cuda::getCurrentCUDAStream().stream()>>>(a.data_ptr<float>(),b.data_ptr<float>(),out.data_ptr<float>(),a.size(1),b.size(2),a.size(2),a.size(0),streams,stride);
    C10_CUDA_KERNEL_LAUNCH_CHECK();return out;
}

torch::Tensor dense_fma(torch::Tensor a,torch::Tensor w,int64_t partitions=1) {
    TORCH_CHECK(a.is_cuda() && w.device()==a.device() && a.dim()==2 && w.dim()==2 && a.size(1)==w.size(1),"Matching CUDA dense matrices required");
    TORCH_CHECK(a.scalar_type()==torch::kFloat32 && w.scalar_type()==torch::kFloat32 && a.is_contiguous() && w.is_contiguous(),"Contiguous FP32 matrices required");
    TORCH_CHECK(partitions>0 && a.size(1)%partitions==0,"Complete contiguous K partitions required");
    c10::cuda::CUDAGuard guard(a.device());auto out=torch::empty({a.size(0),w.size(0)},a.options());
    if(out.numel())dense_fma_kernel<<<(out.numel()+127)/128,128,0,c10::cuda::getCurrentCUDAStream().stream()>>>(a.data_ptr<float>(),w.data_ptr<float>(),out.data_ptr<float>(),a.size(0),w.size(0),a.size(1),partitions);
    C10_CUDA_KERNEL_LAUNCH_CHECK();return out;
}

torch::Tensor mma(torch::Tensor a,torch::Tensor b,c10::optional<torch::Tensor> initial,bool canonical=false,bool integer=true) {
    TORCH_CHECK(a.is_cuda() && b.is_cuda() && a.is_contiguous() && b.is_contiguous(),"Contiguous CUDA matrices required");
    TORCH_CHECK(a.dim()==3 && b.dim()==3 && a.size(0)==b.size(0) && a.size(2)==b.size(1),"Invalid batched matrix shapes");
    TORCH_CHECK(a.scalar_type()==torch::kBFloat16 && b.scalar_type()==torch::kBFloat16 && a.device()==b.device(),"Matching CUDA BF16 required");
    int m=a.size(1),k=a.size(2),n=b.size(2);
    TORCH_CHECK(m%16==0 && n%8==0 && k%16==0 && m>0 && n>0 && k>0,"Complete 16x8x16 tiles required");
    c10::cuda::CUDAGuard guard(a.device());
    auto out=torch::empty({a.size(0),m,n},a.options().dtype(torch::kFloat32));
    const float* cp=nullptr;
    if(initial.has_value()) {
        auto c=*initial;
        TORCH_CHECK(c.device()==a.device() && c.scalar_type()==torch::kFloat32 && c.is_contiguous() && c.sizes()==out.sizes(),"Invalid FP32 accumulator");
        cp=c.data_ptr<float>();
    }
    if(canonical && integer)mma_integer_kernel<<<(a.size(0)*m*n+127)/128,128,0,c10::cuda::getCurrentCUDAStream().stream()>>>(
        reinterpret_cast<const unsigned short*>(a.data_ptr()),reinterpret_cast<const unsigned short*>(b.data_ptr()),cp,out.data_ptr<float>(),m,n,k,a.size(0));
    else if(canonical)mma_canonical_kernel<<<(a.size(0)*m*n+127)/128,128,0,c10::cuda::getCurrentCUDAStream().stream()>>>(
        reinterpret_cast<const __nv_bfloat16*>(a.data_ptr()),reinterpret_cast<const __nv_bfloat16*>(b.data_ptr()),cp,out.data_ptr<float>(),m,n,k,a.size(0));
    else mma_kernel<<<dim3(m/16,n/8,a.size(0)),32,0,c10::cuda::getCurrentCUDAStream().stream()>>>(
        reinterpret_cast<const __nv_bfloat16*>(a.data_ptr()),reinterpret_cast<const __nv_bfloat16*>(b.data_ptr()),cp,out.data_ptr<float>(),m,n,k);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
__global__ void fma_kernel(const float* a,const float* b,const float* c,float* y,int64_t n) {
    int64_t i=int64_t(blockIdx.x)*blockDim.x+threadIdx.x;
    if(i<n) {
        float value;
        asm("fma.rn.f32 %0, %1, %2, %3;":"=f"(value):"f"(a[i]),"f"(b[i]),"f"(c[i]));
        y[i]=value;
    }
}
torch::Tensor tensor_fma(torch::Tensor a,torch::Tensor b,torch::Tensor c) {
    for(auto x:{a,b,c})TORCH_CHECK(x.is_cuda() && x.device()==a.device() && x.scalar_type()==torch::kFloat32 && x.is_contiguous() && x.sizes()==a.sizes(),"Matching contiguous CUDA FP32 operands required");
    c10::cuda::CUDAGuard guard(a.device());auto y=torch::empty_like(a);
    if(a.numel())fma_kernel<<<(a.numel()+255)/256,256,0,c10::cuda::getCurrentCUDAStream().stream()>>>(a.data_ptr<float>(),b.data_ptr<float>(),c.data_ptr<float>(),y.data_ptr<float>(),a.numel());
    C10_CUDA_KERNEL_LAUNCH_CHECK();return y;
}
__global__ void bc_norm_kernel(const __nv_bfloat16* x,const float* weight,__nv_bfloat16* output,float* inverse,int rows,float eps) {
    int row=blockIdx.x,lane=threadIdx.x;
    float a=__bfloat162float(x[row*64+lane*2]),b=__bfloat162float(x[row*64+lane*2+1]);
    float second=__fmul_rn(b,b),sum=__fmaf_rn(a,a,second);
    #pragma unroll
    for(int offset=16;offset;offset>>=1)sum=__fadd_rn(sum,__shfl_xor_sync(0xffffffff,sum,offset));
    float mean,root,reciprocal;
    asm("div.full.f32 %0, %1, %2;":"=f"(mean):"f"(sum),"f"(64.f));
    mean=__fadd_rn(mean,eps);
    asm("sqrt.approx.ftz.f32 %0, %1;":"=f"(root):"f"(mean));
    asm("div.full.f32 %0, %1, %2;":"=f"(reciprocal):"f"(1.f),"f"(root));
    if(lane==0)inverse[row]=reciprocal;
    output[row*64+lane*2]=__float2bfloat16_rn(__fmul_rn(__fmul_rn(a,reciprocal),weight[lane*2]));
    output[row*64+lane*2+1]=__float2bfloat16_rn(__fmul_rn(__fmul_rn(b,reciprocal),weight[lane*2+1]));
}
std::vector<torch::Tensor> bc_norm(torch::Tensor x,torch::Tensor weight,double eps) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type()==torch::kBFloat16 && x.is_contiguous() && x.size(-1)==64,"Contiguous CUDA BF16 state64 required");
    TORCH_CHECK(weight.is_cuda() && weight.device()==x.device() && weight.scalar_type()==torch::kFloat32 && weight.is_contiguous() && weight.numel()==64,"CUDA FP32 norm weight required");
    c10::cuda::CUDAGuard guard(x.device());
    auto output=torch::empty_like(x),inverse=torch::empty({x.numel()/64},x.options().dtype(torch::kFloat32));
    if(x.numel())bc_norm_kernel<<<x.numel()/64,32,0,c10::cuda::getCurrentCUDAStream().stream()>>>(
        reinterpret_cast<__nv_bfloat16*>(x.data_ptr()),weight.data_ptr<float>(),reinterpret_cast<__nv_bfloat16*>(output.data_ptr()),inverse.data_ptr<float>(),x.numel()/64,float(eps));
    C10_CUDA_KERNEL_LAUNCH_CHECK();return {output,inverse};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME,m) {m.def("unary",&unary);m.def("bmm_reference",&bmm_reference);m.def("mma",&mma,pybind11::arg("a"),pybind11::arg("b"),pybind11::arg("initial"),pybind11::arg("canonical")=false,pybind11::arg("integer")=true);m.def("dense_fma",&dense_fma,pybind11::arg("a"),pybind11::arg("w"),pybind11::arg("partitions")=1);m.def("fma",&tensor_fma);m.def("bc_norm",&bc_norm);}