#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <math.h>

#define EPSILON 1e-6f

// Optimized L1 normalization kernel using unrolled loops and warp-level reduction.
// Each block processes one row (of shape [dim]) of the input tensor.
// First, each thread calculates a partial sum of the absolute values with the loop explicitly unrolled.
// Then, a shared-memory reduction is performed (with unrolling) and the final warp is reduced explicitly.
// Finally, each element is normalized by the computed L1 norm.
__global__ void l1norm_kernel(const float* __restrict__ input,
                              float* __restrict__ output,
                              int dim) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    
    extern __shared__ float sdata[];  // shared memory for reduction
    
    const float* row_in = input + row * dim;
    float sum = 0.0f;
    
    // Compute partial sum over elements; unroll loop for performance.
    #pragma unroll
    for (int col = tid; col < dim; col += blockDim.x) {
        sum += fabsf(row_in[col]);
    }
    sdata[tid] = sum;
    __syncthreads();
    
    // Shared memory reduction: reduce in powers of two until 32 threads remain.
    for (unsigned int s = blockDim.x >> 1; s >= 32; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }
    
    // Warp-level reduction with explicit unrolling.
    if (tid < 32) {
        float sum_val = sdata[tid];
        sum_val += __shfl_down_sync(0xffffffff, sum_val, 16);
        sum_val += __shfl_down_sync(0xffffffff, sum_val, 8);
        sum_val += __shfl_down_sync(0xffffffff, sum_val, 4);
        sum_val += __shfl_down_sync(0xffffffff, sum_val, 2);
        sum_val += __shfl_down_sync(0xffffffff, sum_val, 1);
        if (tid == 0) {
            sdata[0] = sum_val;
        }
    }
    __syncthreads();
    
    float l1 = sdata[0] + EPSILON;
    
    // Normalize each element in the row.
    #pragma unroll
    for (int col = tid; col < dim; col += blockDim.x) {
        output[row * dim + col] = row_in[col] / l1;
    }
}

torch::Tensor l1norm(torch::Tensor input) {
    // Ensure the input tensor is on the correct CUDA device.
    c10::cuda::CUDAGuard device_guard(input.device());
    
    int batch = input.size(0);
    int dim = input.size(1);
    
    auto output = torch::empty_like(input);
    
    // Dynamically determine thread block size.
    // Use 512 threads if the dimension is >= 512; otherwise, choose the next power-of-two >= dim.
    int threads = 512;
    if(dim < threads) {
        threads = 1;
        while(threads < dim) {
            threads *= 2;
        }
    }
    
    dim3 block(threads);
    dim3 grid(batch);
    size_t shared_mem = threads * sizeof(float);
    
    l1norm_kernel<<<grid, block, shared_mem, c10::cuda::getCurrentCUDAStream()>>>(
        input.data_ptr<float>(), output.data_ptr<float>(), dim);
    
    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("l1norm", &l1norm, "Optimized fused L1 normalization kernel with dynamic thread block sizing (CUDA)");
}