#ifndef __CAUSAL_SOFTMAX_KERNEL_CUH__
#define __CAUSAL_SOFTMAX_KERNEL_CUH__

template <unsigned int BLOCK_SIZE, typename Tdata, typename Tcompute>
__device__ void causalSoftmaxKernel(
    Tdata *y_, const Tdata *x_,
    size_t batch, size_t height, size_t width,
    ptrdiff_t y_stride_b, ptrdiff_t y_stride_h,
    ptrdiff_t x_stride_b, ptrdiff_t x_stride_h) {

    Tdata *y = y_                       // threadIdx.x for col_id
             + blockIdx.y * y_stride_b  // gridDim.y for batch_id
             + blockIdx.x * y_stride_h; // gridDim.x for row_id
    const Tdata *x = x_ + blockIdx.y * x_stride_b + blockIdx.x * x_stride_h;

    // [Reduce] Find max value in each row and store in shared memory
    __shared__ Tdata max_;
    Tdata max_0 = op::common_cuda::reduce_op::max<BLOCK_SIZE, Tdata>(x, width - height + 1 + blockIdx.x);
    if (threadIdx.x == 0) {
        max_ = max_0;
    }
    __syncthreads();

    // [Elementwise] Subtract max value from each element and apply causal mask
    for (size_t col = threadIdx.x; col < width; col += BLOCK_SIZE) {
        //   row_id ↓ |<-     width   ->|
        //          0 | * * * ... *     |
        //          1 | * * * ... * *   |
        //          2 | * * * ... * * * |
        //  height: 3  col_id->
        if (width + blockIdx.x >= col + height) {
            if constexpr (std::is_same_v<Tdata, half> || std::is_same_v<Tdata, cuda_bfloat16>) {
                /*
                 * MUSA does not support CUDA's native `hexp` function.
                 * This code performs an explicit conversion:
                 * it casts the input to `float`, computes the exponential, and casts the result back.
                 * This ensures compatibility and correct behavior on the MUSA platform.
                 */
                float val = static_cast<float>(x[col]) - static_cast<float>(max_);
                y[col] = static_cast<Tdata>(expf(val));
            } else {
                y[col] = exp(x[col] - max_);
            }
        } else {
            /*
             * In MUSA, the `__mt_bfloat16` type has ambiguous constructors for integer literals (e.g., `0`),
             * as it could be implicitly converted from either `float` or `double`.
             *
             * This differs from CUDA's `half` type, which can typically be initialized
             * from integer literals without ambiguity.
             *
             * To resolve this, we use the float literal `0.0f` to explicitly
             * specify the conversion path, ensuring platform compatibility.
             */
            y[col] = Tdata(0.0f);
        }
    }
    __syncthreads();

    // [Reduce] Find the sum of each updated row and store in shared memory
    __shared__ Tcompute sum_;
    Tcompute sum_0 = op::common_cuda::reduce_op::sum<BLOCK_SIZE, Tdata, Tcompute>(y, width);
    if (threadIdx.x == 0) {
        sum_ = sum_0;
    }
    __syncthreads();

    // [Elementwise] Divide each element by the sum and store in shared memory
    for (size_t col = threadIdx.x; col < width; col += BLOCK_SIZE) {
        /*
         * MUSA's bfloat16 type does not have a viable overloaded `/=` operator for float division.
         * This change explicitly casts both operands to `float` before division,
         * and then casts the result back to `Tdata`.
         * This ensures the operation is performed correctly and avoids compilation errors.
         */
        y[col] = static_cast<Tdata>(static_cast<float>(y[col]) / static_cast<float>(sum_));
    }
}

#endif // __CAUSAL_SOFTMAX_KERNEL_CUH__
