metakernel v1: device dossier, throttle-rejected bench, ULP compare, fuzz, sweep, stamps
4611d12 verified | // Local development binding (torch.utils.cpp_extension JIT; Windows/MSVC). | |
| // Mirrors the torch-ext binding's argument order exactly. | |
| namespace { | |
| void chk(bool ok, const char* msg) { TORCH_CHECK(ok, msg); } | |
| cudaStream_t stream_of(const torch::Tensor& t) { | |
| const at::cuda::CUDAGuard guard(t.device()); | |
| return at::cuda::getCurrentCUDAStream(); | |
| } | |
| void mk_triad(torch::Tensor a, torch::Tensor b, torch::Tensor c, | |
| int64_t width, double s) { | |
| chk(a.is_cuda() && a.is_contiguous() && a.dtype() == torch::kFloat32, | |
| "a must be contiguous CUDA f32"); | |
| chk(b.sizes() == a.sizes() && c.sizes() == a.sizes(), "size mismatch"); | |
| chk(width == 1 || width == 2 || width == 4, "width must be 1|2|4"); | |
| chk(a.numel() % width == 0, "numel must divide width"); | |
| mk_triad_launch(a.const_data_ptr<float>(), b.const_data_ptr<float>(), | |
| c.data_ptr<float>(), a.numel() / width, (int)width, | |
| (float)s, stream_of(a)); | |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); | |
| } | |
| void mk_triad_passes(torch::Tensor a, torch::Tensor b, torch::Tensor c, | |
| double s, int64_t passes) { | |
| chk(a.is_cuda() && a.is_contiguous() && a.dtype() == torch::kFloat32, | |
| "a must be contiguous CUDA f32"); | |
| chk(a.numel() % 4 == 0, "numel must divide 4"); | |
| mk_triad_passes_launch(a.const_data_ptr<float>(), b.const_data_ptr<float>(), | |
| c.data_ptr<float>(), a.numel() / 4, 4, (float)s, | |
| (int)passes, stream_of(a)); | |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); | |
| } | |
| void mk_chase_global(torch::Tensor ring, int64_t hops, torch::Tensor out) { | |
| chk(ring.is_cuda() && ring.is_contiguous() && | |
| ring.dtype() == torch::kInt32, | |
| "ring must be contiguous CUDA i32"); | |
| chk(out.is_cuda() && out.dtype() == torch::kInt64 && out.numel() >= 2, | |
| "out must be CUDA i64 [>=2]"); | |
| mk_chase_global_launch(ring.const_data_ptr<int>(), hops, | |
| reinterpret_cast<long long*>(out.data_ptr<int64_t>()), stream_of(ring)); | |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); | |
| } | |
| void mk_read(torch::Tensor a, torch::Tensor sink) { | |
| chk(a.is_cuda() && a.is_contiguous() && a.dtype() == torch::kFloat32 && | |
| a.numel() % 4 == 0, | |
| "a must be contiguous CUDA f32, numel % 4 == 0"); | |
| mk_read_launch(a.const_data_ptr<float>(), a.numel() / 4, | |
| sink.data_ptr<float>(), stream_of(a)); | |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); | |
| } | |
| void mk_write(torch::Tensor c, double v) { | |
| chk(c.is_cuda() && c.is_contiguous() && c.dtype() == torch::kFloat32 && | |
| c.numel() % 4 == 0, | |
| "c must be contiguous CUDA f32, numel % 4 == 0"); | |
| mk_write_launch(c.data_ptr<float>(), c.numel() / 4, (float)v, | |
| stream_of(c)); | |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); | |
| } | |
| void mk_gather(torch::Tensor a, torch::Tensor idx, torch::Tensor c, | |
| int64_t threads, int64_t dyn_smem) { | |
| chk(a.is_cuda() && a.is_contiguous() && a.dtype() == torch::kFloat32 && | |
| a.numel() % 4 == 0, | |
| "a must be contiguous CUDA f32, numel % 4 == 0"); | |
| chk(idx.is_cuda() && idx.dtype() == torch::kInt32 && | |
| idx.numel() * 4 == c.numel() && c.numel() % 4 == 0, | |
| "idx i32 [n4] with c f32 [n4*4]"); | |
| mk_gather_launch(a.const_data_ptr<float>(), idx.const_data_ptr<int>(), | |
| c.data_ptr<float>(), idx.numel(), (int)threads, | |
| (int)dyn_smem, stream_of(a)); | |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); | |
| } | |
| void mk_fma_f64(int64_t blocks, int64_t threads, int64_t iters, | |
| torch::Tensor sink) { | |
| chk(sink.is_cuda() && sink.dtype() == torch::kFloat64 && | |
| sink.numel() >= 1, | |
| "sink must be CUDA f64 [>=1]"); | |
| mk_fma_f64_launch((int)blocks, (int)threads, iters, | |
| sink.data_ptr<double>(), stream_of(sink)); | |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); | |
| } | |
| void mk_fma_dep(int64_t blocks, int64_t threads, int64_t dyn_smem, | |
| int64_t iters, torch::Tensor sink) { | |
| chk(sink.is_cuda() && sink.dtype() == torch::kFloat32 && sink.numel() >= 1, | |
| "sink must be CUDA f32 [>=1]"); | |
| mk_fma_dep_launch((int)blocks, (int)threads, (int)dyn_smem, iters, | |
| sink.data_ptr<float>(), stream_of(sink)); | |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); | |
| } | |
| void mk_chase_shared(torch::Tensor ring, int64_t hops, torch::Tensor out) { | |
| chk(ring.is_cuda() && ring.is_contiguous() && | |
| ring.dtype() == torch::kInt32, | |
| "ring must be contiguous CUDA i32"); | |
| chk(ring.numel() <= 12288, "shared ring must fit 48 KB (<= 12288 ints)"); | |
| chk(out.is_cuda() && out.dtype() == torch::kInt64 && out.numel() >= 2, | |
| "out must be CUDA i64 [>=2]"); | |
| mk_chase_shared_launch(ring.const_data_ptr<int>(), (int)ring.numel(), | |
| hops, reinterpret_cast<long long*>(out.data_ptr<int64_t>()), stream_of(ring)); | |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); | |
| } | |
| void mk_fma_f32(int64_t blocks, int64_t threads, int64_t iters, | |
| torch::Tensor sink) { | |
| chk(sink.is_cuda() && sink.dtype() == torch::kFloat32 && sink.numel() >= 1, | |
| "sink must be CUDA f32 [>=1]"); | |
| mk_fma_f32_launch((int)blocks, (int)threads, iters, | |
| sink.data_ptr<float>(), stream_of(sink)); | |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); | |
| } | |
| int64_t mk_mma(int64_t kind, int64_t blocks, int64_t warps, int64_t iters, | |
| torch::Tensor sink) { | |
| chk(sink.is_cuda() && sink.dtype() == torch::kFloat32 && sink.numel() >= 1, | |
| "sink must be CUDA f32 [>=1]"); | |
| chk(kind >= 0 && kind <= 4, "kind must be 0..4"); | |
| cudaDeviceProp prop; | |
| cudaGetDeviceProperties(&prop, sink.device().index()); | |
| int r = mk_mma_launch((int)kind, (int)blocks, (int)warps, iters, | |
| sink.data_ptr<float>(), prop.major, prop.minor, | |
| stream_of(sink)); | |
| if (r == 0) C10_CUDA_KERNEL_LAUNCH_CHECK(); | |
| return r; | |
| } | |
| void mk_atomics(int64_t mode, int64_t blocks, int64_t threads, int64_t iters, | |
| torch::Tensor buf, int64_t slots) { | |
| chk(buf.is_cuda() && buf.dtype() == torch::kFloat32 && buf.numel() >= 2, | |
| "buf must be CUDA f32 [>=2]"); | |
| chk(slots >= 1 && slots <= buf.numel(), "slots in [1, buf.numel()]"); | |
| mk_atomics_launch((int)mode, (int)blocks, (int)threads, iters, | |
| buf.data_ptr<float>(), slots, stream_of(buf)); | |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); | |
| } | |
| void mk_spin(int64_t blocks, int64_t threads, int64_t ticks, | |
| torch::Tensor out, torch::Tensor sink) { | |
| chk(out.is_cuda() && out.dtype() == torch::kInt64 && out.numel() >= 1, | |
| "out must be CUDA i64 [>=1]"); | |
| mk_spin_launch((int)blocks, (int)threads, ticks, | |
| reinterpret_cast<long long*>(out.data_ptr<int64_t>()), sink.data_ptr<float>(), | |
| stream_of(out)); | |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); | |
| } | |
| void mk_empty(torch::Tensor device_anchor) { | |
| mk_empty_launch(stream_of(device_anchor)); | |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); | |
| } | |
| int64_t mk_barrier(int64_t blocks, int64_t threads, int64_t iters, | |
| torch::Tensor out) { | |
| chk(out.is_cuda() && out.dtype() == torch::kInt64 && out.numel() >= 1, | |
| "out must be CUDA i64 [>=1]"); | |
| int r = mk_barrier_launch((int)blocks, (int)threads, (int)iters, | |
| reinterpret_cast<long long*>(out.data_ptr<int64_t>()), stream_of(out)); | |
| if (r == 0) C10_CUDA_KERNEL_LAUNCH_CHECK(); | |
| return r; | |
| } | |
| int64_t mk_occupancy(int64_t threads, int64_t dyn_smem, | |
| torch::Tensor device_anchor) { | |
| const at::cuda::CUDAGuard guard(device_anchor.device()); | |
| int blocks_per_sm = -1; | |
| int r = mk_occupancy_triad((int)threads, (int)dyn_smem, &blocks_per_sm); | |
| return r == 0 ? blocks_per_sm : -1; | |
| } | |
| } // namespace | |
| PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { | |
| m.def("mk_triad", &mk_triad); | |
| m.def("mk_triad_passes", &mk_triad_passes); | |
| m.def("mk_read", &mk_read); | |
| m.def("mk_write", &mk_write); | |
| m.def("mk_gather", &mk_gather); | |
| m.def("mk_fma_f64", &mk_fma_f64); | |
| m.def("mk_fma_dep", &mk_fma_dep); | |
| m.def("mk_chase_global", &mk_chase_global); | |
| m.def("mk_chase_shared", &mk_chase_shared); | |
| m.def("mk_fma_f32", &mk_fma_f32); | |
| m.def("mk_mma", &mk_mma); | |
| m.def("mk_atomics", &mk_atomics); | |
| m.def("mk_spin", &mk_spin); | |
| m.def("mk_empty", &mk_empty); | |
| m.def("mk_barrier", &mk_barrier); | |
| m.def("mk_occupancy", &mk_occupancy); | |
| } | |