AMConv2d_grad_cuda_v2.cu

#include <torch/extension.h>

#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
#include <cmath>

    // CUDA kernel for the convolution operation
    template <typename scalar_t>
    __global__ void amconv2d_LUTgrad_kernel(
        const torch::PackedTensorAccessor32<scalar_t,4,torch::RestrictPtrTraits> input,
        const torch::PackedTensorAccessor32<scalar_t,4,torch::RestrictPtrTraits> weights,
        torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> output,
        const torch::PackedTensorAccessor32<scalar_t,4,torch::RestrictPtrTraits> grad,
        int batch_size,
        int in_channels,
        int input_height,
        int input_width,
        int out_channels,
        int kernel_height,
        int kernel_width,
        int output_height,
        int output_width,
        int stride_h,
        int stride_w,
        float input_scale,
        float input_offset,
        const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> weight_scale,
        const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> weight_offset,
        const int row_dim,
        const int col_dim
        ) {
        
        int out_idx = blockIdx.x * blockDim.x + threadIdx.x;
        int total_elements = batch_size * out_channels * output_height * output_width;
 


        if (out_idx < total_elements) {
            int b = out_idx / (out_channels * output_height * output_width); //第几个batch
            int oc = (out_idx / (output_height * output_width)) % out_channels;//第几个channel
            int oh = (out_idx / output_width) % output_height;//第几个height
            int ow = out_idx % output_width;//第几个width

            int h_start = oh * stride_h;
            int w_start = ow * stride_w;

            for (int ic = 0; ic < in_channels; ++ic) {
                for (int kh = 0; kh < kernel_height; ++kh) {
                    for (int kw = 0; kw < kernel_width; ++kw) {
                        int h = h_start + kh;
                        int w = w_start + kw;

                        if (h < input_height && w < input_width) {
                            float input_value = input[b][ic][h][w];
                            float weight_value = weights[oc][ic][kh][kw];
                            int quant_input_value=std::nearbyint((input_value-input_offset)/input_scale);
                            int quant_weight_value=std::nearbyint((weight_value-weight_offset[oc][ic])/weight_scale[oc][ic]);

                            atomicAdd(&output[quant_input_value][quant_weight_value], input_scale*weight_scale[oc][ic]*grad[b][oc][oh][ow]);
                        }
                    }
                }
            }
        }

        


    }

    // Forward pass function that launches the CUDA kernel
    torch::Tensor AMConv2d_LUTgrad(
        torch::Tensor input,
        torch::Tensor weights,
        torch::Tensor stride,
        torch::Tensor grad,
        int row_dim,
        int col_dim,
        float input_scale,
        float input_offset,
        torch::Tensor weight_scale,
        torch::Tensor weight_offset) {

        // Ensure inputs are of the correct type
        TORCH_CHECK(input.dim() == 4, "Input should be a 4D tensor");
        TORCH_CHECK(weights.dim() == 4, "Weights should be a 4D tensor");

        int stride_h = stride[0].item<int>();
        int stride_w = stride[1].item<int>();
        
        // Get input dimensions
        auto batch_size = input.size(0);
        auto in_channels = input.size(1);
        auto input_height = input.size(2);
        auto input_width = input.size(3);

        // Get weight dimensions
        auto out_channels = weights.size(0);
        auto kernel_height = weights.size(2);
        auto kernel_width = weights.size(3);

        // Calculate output dimensions
        int output_height = grad.size(2);
        int output_width = grad.size(3);


        auto output = torch::zeros({row_dim,col_dim}, input.options());

        int total_elements = batch_size * out_channels * output_height * output_width;
        int threads = 1024;
        int blocks = (total_elements + threads - 1) / threads;
        
        // Launch the CUDA kernel
        AT_DISPATCH_FLOATING_TYPES(input.type(), "amconv2d_grad_cuda", ([&] {
        amconv2d_LUTgrad_kernel<scalar_t><<<blocks, threads>>>(
            input.packed_accessor32<scalar_t,4,torch::RestrictPtrTraits>(),
            weights.packed_accessor32<scalar_t,4,torch::RestrictPtrTraits>(),
            output.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
            grad.packed_accessor32<scalar_t,4,torch::RestrictPtrTraits>(),
            batch_size, 
            in_channels,
            input_height,
            input_width,
            out_channels,
            kernel_height,
            kernel_width,
            output_height,
            output_width,
            stride_h,
            stride_w,
            input_scale,
            input_offset,
            weight_scale.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
            weight_offset.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
            row_dim,
            col_dim);}));

        return {output};
    }



PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("LUTgrad", &AMConv2d_LUTgrad, "AMConv2d forward pass (CUDA)");
}