submission 778362

import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t


CPP_SRC = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

#include <stdexcept>
#include <string>

static const char* kPtx = R"ptx(
.version 6.0
.target sm_60
.address_size 64

.visible .entry block_sum_kernel_float(
	.param .u64 block_sum_kernel_float_param_0,
	.param .u64 block_sum_kernel_float_param_1,
	.param .u32 block_sum_kernel_float_param_2
)
{
	.reg .pred 	%p<20>;
	.reg .f32 	%f<11>;
	.reg .b32 	%r<37>;
	.reg .f64 	%fd<70>;
	.reg .b64 	%rd<23>;
	.shared .align 8 .b8 _ZZ22block_sum_kernel_floatE4smem[2048];

	ld.param.u64 	%rd15, [block_sum_kernel_float_param_0];
	ld.param.u64 	%rd14, [block_sum_kernel_float_param_1];
	ld.param.u32 	%r18, [block_sum_kernel_float_param_2];
	cvta.to.global.u64 	%rd1, %rd15;
	mov.u32 	%r19, %ctaid.x;
	shl.b32 	%r1, %r19, 9;
	mov.u32 	%r2, %tid.x;
	add.s32 	%r35, %r1, %r2;
	mov.u32 	%r20, %nctaid.x;
	shl.b32 	%r4, %r20, 9;
	setp.ge.s32 	%p1, %r35, %r18;
	mov.f64 	%fd64, 0d0000000000000000;
	@%p1 bra 	$L__BB0_17;

	not.b32 	%r21, %r2;
	add.s32 	%r22, %r21, %r18;
	sub.s32 	%r23, %r22, %r1;
	div.u32 	%r5, %r23, %r4;
	add.s32 	%r24, %r5, 1;
	and.b32  	%r34, %r24, 3;
	setp.eq.s32 	%p2, %r34, 0;
	mov.f64 	%fd64, 0d0000000000000000;
	@%p2 bra 	$L__BB0_6;

	mul.wide.s32 	%rd16, %r35, 4;
	add.s64 	%rd17, %rd1, %rd16;
	add.s64 	%rd22, %rd17, 1024;
	mul.wide.s32 	%rd3, %r4, 4;
	mov.f64 	%fd64, 0d0000000000000000;

$L__BB0_3:
	.pragma "nounroll";
	add.s64 	%rd5, %rd22, -1024;
	ld.global.nc.f32 	%f1, [%rd22+-1024];
	cvt.f64.f32 	%fd25, %f1;
	add.f64 	%fd64, %fd64, %fd25;
	add.s32 	%r25, %r35, 256;
	setp.ge.s32 	%p3, %r25, %r18;
	@%p3 bra 	$L__BB0_5;

	ld.global.nc.f32 	%f2, [%rd5+1024];
	cvt.f64.f32 	%fd26, %f2;
	add.f64 	%fd64, %fd64, %fd26;

$L__BB0_5:
	add.s32 	%r35, %r35, %r4;
	add.s64 	%rd22, %rd22, %rd3;
	add.s32 	%r34, %r34, -1;
	setp.ne.s32 	%p4, %r34, 0;
	@%p4 bra 	$L__BB0_3;

$L__BB0_6:
	setp.lt.u32 	%p5, %r5, 3;
	@%p5 bra 	$L__BB0_17;

	mul.wide.s32 	%rd7, %r4, 4;

$L__BB0_8:
	mul.wide.s32 	%rd18, %r35, 4;
	add.s64 	%rd8, %rd1, %rd18;
	ld.global.nc.f32 	%f3, [%rd8];
	cvt.f64.f32 	%fd27, %f3;
	add.f64 	%fd65, %fd64, %fd27;
	add.s32 	%r26, %r35, 256;
	setp.ge.s32 	%p6, %r26, %r18;
	@%p6 bra 	$L__BB0_10;

	ld.global.nc.f32 	%f4, [%rd8+1024];
	cvt.f64.f32 	%fd28, %f4;
	add.f64 	%fd65, %fd65, %fd28;

$L__BB0_10:
	add.s64 	%rd10, %rd8, %rd7;
	ld.global.nc.f32 	%f5, [%rd10];
	cvt.f64.f32 	%fd29, %f5;
	add.f64 	%fd66, %fd65, %fd29;
	add.s32 	%r13, %r35, %r4;
	add.s32 	%r27, %r13, 256;
	setp.ge.s32 	%p7, %r27, %r18;
	@%p7 bra 	$L__BB0_12;

	ld.global.nc.f32 	%f6, [%rd10+1024];
	cvt.f64.f32 	%fd30, %f6;
	add.f64 	%fd66, %fd66, %fd30;

$L__BB0_12:
	add.s64 	%rd12, %rd10, %rd7;
	ld.global.nc.f32 	%f7, [%rd12];
	cvt.f64.f32 	%fd31, %f7;
	add.f64 	%fd67, %fd66, %fd31;
	add.s32 	%r14, %r13, %r4;
	add.s32 	%r28, %r14, 256;
	setp.ge.s32 	%p8, %r28, %r18;
	@%p8 bra 	$L__BB0_14;

	ld.global.nc.f32 	%f8, [%rd12+1024];
	cvt.f64.f32 	%fd32, %f8;
	add.f64 	%fd67, %fd67, %fd32;

$L__BB0_14:
	add.s64 	%rd13, %rd12, %rd7;
	ld.global.nc.f32 	%f9, [%rd13];
	cvt.f64.f32 	%fd33, %f9;
	add.f64 	%fd64, %fd67, %fd33;
	add.s32 	%r15, %r14, %r4;
	add.s32 	%r29, %r15, 256;
	setp.ge.s32 	%p9, %r29, %r18;
	@%p9 bra 	$L__BB0_16;

	ld.global.nc.f32 	%f10, [%rd13+1024];
	cvt.f64.f32 	%fd34, %f10;
	add.f64 	%fd64, %fd64, %fd34;

$L__BB0_16:
	add.s32 	%r35, %r15, %r4;
	setp.lt.s32 	%p10, %r35, %r18;
	@%p10 bra 	$L__BB0_8;

$L__BB0_17:
	shl.b32 	%r30, %r2, 3;
	mov.u32 	%r31, _ZZ22block_sum_kernel_floatE4smem;
	add.s32 	%r17, %r31, %r30;
	st.shared.f64 	[%r17], %fd64;
	bar.sync 	0;
	setp.gt.s32 	%p11, %r2, 127;
	@%p11 bra 	$L__BB0_19;

	ld.shared.f64 	%fd35, [%r17];
	ld.shared.f64 	%fd36, [%r17+1024];
	add.f64 	%fd37, %fd36, %fd35;
	st.shared.f64 	[%r17], %fd37;

$L__BB0_19:
	bar.sync 	0;
	setp.gt.s32 	%p12, %r2, 63;
	@%p12 bra 	$L__BB0_21;

	ld.shared.f64 	%fd38, [%r17];
	ld.shared.f64 	%fd39, [%r17+512];
	add.f64 	%fd40, %fd39, %fd38;
	st.shared.f64 	[%r17], %fd40;

$L__BB0_21:
	bar.sync 	0;
	setp.gt.s32 	%p13, %r2, 31;
	@%p13 bra 	$L__BB0_23;

	ld.shared.f64 	%fd41, [%r17];
	ld.shared.f64 	%fd42, [%r17+256];
	add.f64 	%fd43, %fd42, %fd41;
	st.shared.f64 	[%r17], %fd43;

$L__BB0_23:
	bar.sync 	0;
	setp.gt.s32 	%p14, %r2, 15;
	@%p14 bra 	$L__BB0_25;

	ld.shared.f64 	%fd44, [%r17];
	ld.shared.f64 	%fd45, [%r17+128];
	add.f64 	%fd46, %fd45, %fd44;
	st.shared.f64 	[%r17], %fd46;

$L__BB0_25:
	bar.sync 	0;
	setp.gt.s32 	%p15, %r2, 7;
	@%p15 bra 	$L__BB0_27;

	ld.shared.f64 	%fd47, [%r17];
	ld.shared.f64 	%fd48, [%r17+64];
	add.f64 	%fd49, %fd48, %fd47;
	st.shared.f64 	[%r17], %fd49;

$L__BB0_27:
	bar.sync 	0;
	setp.gt.s32 	%p16, %r2, 3;
	@%p16 bra 	$L__BB0_29;

	ld.shared.f64 	%fd50, [%r17];
	ld.shared.f64 	%fd51, [%r17+32];
	add.f64 	%fd52, %fd51, %fd50;
	st.shared.f64 	[%r17], %fd52;

$L__BB0_29:
	bar.sync 	0;
	setp.gt.s32 	%p17, %r2, 1;
	@%p17 bra 	$L__BB0_31;

	ld.shared.f64 	%fd53, [%r17];
	ld.shared.f64 	%fd54, [%r17+16];
	add.f64 	%fd55, %fd54, %fd53;
	st.shared.f64 	[%r17], %fd55;

$L__BB0_31:
	bar.sync 	0;
	setp.gt.s32 	%p18, %r2, 0;
	@%p18 bra 	$L__BB0_33;

	ld.shared.f64 	%fd56, [%r17];
	ld.shared.f64 	%fd57, [%r17+8];
	add.f64 	%fd58, %fd57, %fd56;
	st.shared.f64 	[%r17], %fd58;

$L__BB0_33:
	bar.sync 	0;
	setp.ne.s32 	%p19, %r2, 0;
	@%p19 bra 	$L__BB0_35;

	ld.shared.f64 	%fd59, [_ZZ22block_sum_kernel_floatE4smem];
	cvta.to.global.u64 	%rd19, %rd14;
	mul.wide.u32 	%rd20, %r19, 8;
	add.s64 	%rd21, %rd19, %rd20;
	st.global.f64 	[%rd21], %fd59;

$L__BB0_35:
	ret;
}

.visible .entry final_sum_kernel_float(
	.param .u64 final_sum_kernel_float_param_0,
	.param .u64 final_sum_kernel_float_param_1,
	.param .u32 final_sum_kernel_float_param_2
)
{
	.reg .pred 	%p<15>;
	.reg .f32 	%f<2>;
	.reg .b32 	%r<22>;
	.reg .f64 	%fd<50>;
	.reg .b64 	%rd<16>;
	.shared .align 8 .b8 _ZZ22final_sum_kernel_floatE4smem[2048];

	ld.param.u64 	%rd9, [final_sum_kernel_float_param_0];
	ld.param.u64 	%rd8, [final_sum_kernel_float_param_1];
	ld.param.u32 	%r12, [final_sum_kernel_float_param_2];
	cvta.to.global.u64 	%rd1, %rd9;
	mov.u32 	%r1, %tid.x;
	setp.ge.s32 	%p1, %r1, %r12;
	mov.f64 	%fd49, 0d0000000000000000;
	@%p1 bra 	$L__BB1_7;

	not.b32 	%r13, %r1;
	add.s32 	%r2, %r13, %r12;
	shr.u32 	%r14, %r2, 8;
	add.s32 	%r15, %r14, 1;
	and.b32  	%r19, %r15, 3;
	setp.eq.s32 	%p2, %r19, 0;
	mov.f64 	%fd49, 0d0000000000000000;
	mov.u32 	%r20, %r1;
	@%p2 bra 	$L__BB1_4;

	mul.wide.s32 	%rd10, %r1, 8;
	add.s64 	%rd14, %rd1, %rd10;
	mov.f64 	%fd49, 0d0000000000000000;
	mov.u32 	%r20, %r1;

$L__BB1_3:
	.pragma "nounroll";
	ld.global.nc.f64 	%fd12, [%rd14];
	add.f64 	%fd49, %fd49, %fd12;
	add.s32 	%r20, %r20, 256;
	add.s64 	%rd14, %rd14, 2048;
	add.s32 	%r19, %r19, -1;
	setp.ne.s32 	%p3, %r19, 0;
	@%p3 bra 	$L__BB1_3;

$L__BB1_4:
	setp.lt.u32 	%p4, %r2, 768;
	@%p4 bra 	$L__BB1_7;

	mul.wide.s32 	%rd11, %r20, 8;
	add.s64 	%rd12, %rd1, %rd11;
	add.s64 	%rd15, %rd12, 4096;

$L__BB1_6:
	ld.global.nc.f64 	%fd13, [%rd15+-4096];
	add.f64 	%fd14, %fd49, %fd13;
	ld.global.nc.f64 	%fd15, [%rd15+-2048];
	add.f64 	%fd16, %fd14, %fd15;
	ld.global.nc.f64 	%fd17, [%rd15];
	add.f64 	%fd18, %fd16, %fd17;
	ld.global.nc.f64 	%fd19, [%rd15+2048];
	add.f64 	%fd49, %fd18, %fd19;
	add.s64 	%rd15, %rd15, 8192;
	add.s32 	%r20, %r20, 1024;
	setp.lt.s32 	%p5, %r20, %r12;
	@%p5 bra 	$L__BB1_6;

$L__BB1_7:
	shl.b32 	%r16, %r1, 3;
	mov.u32 	%r17, _ZZ22final_sum_kernel_floatE4smem;
	add.s32 	%r11, %r17, %r16;
	st.shared.f64 	[%r11], %fd49;
	bar.sync 	0;
	setp.gt.s32 	%p6, %r1, 127;
	@%p6 bra 	$L__BB1_9;

	ld.shared.f64 	%fd20, [%r11];
	ld.shared.f64 	%fd21, [%r11+1024];
	add.f64 	%fd22, %fd21, %fd20;
	st.shared.f64 	[%r11], %fd22;

$L__BB1_9:
	bar.sync 	0;
	setp.gt.s32 	%p7, %r1, 63;
	@%p7 bra 	$L__BB1_11;

	ld.shared.f64 	%fd23, [%r11];
	ld.shared.f64 	%fd24, [%r11+512];
	add.f64 	%fd25, %fd24, %fd23;
	st.shared.f64 	[%r11], %fd25;

$L__BB1_11:
	bar.sync 	0;
	setp.gt.s32 	%p8, %r1, 31;
	@%p8 bra 	$L__BB1_13;

	ld.shared.f64 	%fd26, [%r11];
	ld.shared.f64 	%fd27, [%r11+256];
	add.f64 	%fd28, %fd27, %fd26;
	st.shared.f64 	[%r11], %fd28;

$L__BB1_13:
	bar.sync 	0;
	setp.gt.s32 	%p9, %r1, 15;
	@%p9 bra 	$L__BB1_15;

	ld.shared.f64 	%fd29, [%r11];
	ld.shared.f64 	%fd30, [%r11+128];
	add.f64 	%fd31, %fd30, %fd29;
	st.shared.f64 	[%r11], %fd31;

$L__BB1_15:
	bar.sync 	0;
	setp.gt.s32 	%p10, %r1, 7;
	@%p10 bra 	$L__BB1_17;

	ld.shared.f64 	%fd32, [%r11];
	ld.shared.f64 	%fd33, [%r11+64];
	add.f64 	%fd34, %fd33, %fd32;
	st.shared.f64 	[%r11], %fd34;

$L__BB1_17:
	bar.sync 	0;
	setp.gt.s32 	%p11, %r1, 3;
	@%p11 bra 	$L__BB1_19;

	ld.shared.f64 	%fd35, [%r11];
	ld.shared.f64 	%fd36, [%r11+32];
	add.f64 	%fd37, %fd36, %fd35;
	st.shared.f64 	[%r11], %fd37;

$L__BB1_19:
	bar.sync 	0;
	setp.gt.s32 	%p12, %r1, 1;
	@%p12 bra 	$L__BB1_21;

	ld.shared.f64 	%fd38, [%r11];
	ld.shared.f64 	%fd39, [%r11+16];
	add.f64 	%fd40, %fd39, %fd38;
	st.shared.f64 	[%r11], %fd40;

$L__BB1_21:
	bar.sync 	0;
	setp.gt.s32 	%p13, %r1, 0;
	@%p13 bra 	$L__BB1_23;

	ld.shared.f64 	%fd41, [%r11];
	ld.shared.f64 	%fd42, [%r11+8];
	add.f64 	%fd43, %fd42, %fd41;
	st.shared.f64 	[%r11], %fd43;

$L__BB1_23:
	bar.sync 	0;
	setp.ne.s32 	%p14, %r1, 0;
	@%p14 bra 	$L__BB1_25;

	ld.shared.f64 	%fd44, [_ZZ22final_sum_kernel_floatE4smem];
	cvt.rn.f32.f64 	%f1, %fd44;
	cvta.to.global.u64 	%rd13, %rd8;
	st.global.f32 	[%rd13], %f1;

$L__BB1_25:
	ret;
}
)ptx";

static void check_driver(CUresult result, const char* what) {
    if (result != CUDA_SUCCESS) {
        const char* message = nullptr;
        cuGetErrorString(result, &message);
        throw std::runtime_error(std::string(what) + ": " +
                                 (message ? message : "unknown CUDA driver error"));
    }
}

static void check_runtime(cudaError_t result, const char* what) {
    if (result != cudaSuccess) {
        throw std::runtime_error(std::string(what) + ": " +
                                 cudaGetErrorString(result));
    }
}

struct PtxModule {
    CUmodule module = nullptr;
    CUfunction block_sum = nullptr;
    CUfunction final_sum = nullptr;

    PtxModule() {
        check_driver(cuInit(0), "cuInit");
        check_driver(cuModuleLoadData(&module, kPtx), "cuModuleLoadData");
        check_driver(cuModuleGetFunction(&block_sum, module, "block_sum_kernel_float"),
                     "cuModuleGetFunction block_sum_kernel_float");
        check_driver(cuModuleGetFunction(&final_sum, module, "final_sum_kernel_float"),
                     "cuModuleGetFunction final_sum_kernel_float");
    }
};

static PtxModule& ptx_module() {
    static PtxModule instance;
    return instance;
}

torch::Tensor my_op(torch::Tensor input, torch::Tensor output) {
    const int n = static_cast<int>(input.numel());
    constexpr int threads = 256;
    constexpr int max_blocks = 1024;
    int blocks = (n + threads * 2 - 1) / (threads * 2);
    if (blocks > max_blocks) {
        blocks = max_blocks;
    }
    if (blocks < 1) {
        blocks = 1;
    }

    auto partial = torch::empty({blocks}, input.options().dtype(torch::kFloat64));
    CUdeviceptr input_ptr = reinterpret_cast<CUdeviceptr>(input.data_ptr<float>());
    CUdeviceptr partial_ptr = reinterpret_cast<CUdeviceptr>(partial.data_ptr<double>());
    CUdeviceptr output_ptr = reinterpret_cast<CUdeviceptr>(output.data_ptr<float>());

    PtxModule& mod = ptx_module();
    void* block_args[] = {&input_ptr, &partial_ptr, const_cast<int*>(&n)};
    check_driver(cuLaunchKernel(mod.block_sum, blocks, 1, 1, threads, 1, 1, 0, 0,
                                block_args, nullptr),
                 "cuLaunchKernel block_sum_kernel_float");

    void* final_args[] = {&partial_ptr, &output_ptr, &blocks};
    check_driver(cuLaunchKernel(mod.final_sum, 1, 1, 1, threads, 1, 1, 0, 0,
                                final_args, nullptr),
                 "cuLaunchKernel final_sum_kernel_float");

    check_runtime(cudaGetLastError(), "cudaGetLastError");
    return output;
}
"""


module = load_inline(
    name="my_ptx_fast_sum_module",
    cpp_sources=[CPP_SRC],
    functions=["my_op"],
    extra_ldflags=["-lcuda"],
    with_cuda=True,
    verbose=True,
)


def custom_kernel(data: input_t) -> output_t:
    input, output = data
    return module.my_op(input, output)[0]