| #include <torch/library.h> |
| #include <torch/all.h> |
|
|
| #include <ATen/cuda/CUDAContext.h> |
| #include <c10/cuda/CUDAGuard.h> |
|
|
| #include "registration.h" |
| #include "torch_binding.h" |
| #include "../metakernel_cuda/probes_launch.h" |
|
|
| 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"); |
| chk(passes >= 1 && passes <= 1000000, "passes in [1, 1e6]"); |
| 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_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]"); |
| chk(threads >= 32 && threads <= 1024 && dyn_smem >= 0 && |
| dyn_smem <= 48 * 1024, |
| "bad launch shape"); |
| 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]"); |
| chk(dyn_smem >= 0 && dyn_smem <= 48 * 1024, "dyn_smem in [0, 48K]"); |
| 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_fma_ilp(int64_t chains, 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]"); |
| chk(chains == 1 || chains == 2 || chains == 4, "chains must be 1|2|4"); |
| chk(blocks >= 1 && threads >= 32 && threads <= 1024, "bad launch shape"); |
| mk_fma_ilp_launch((int)chains, (int)blocks, (int)threads, iters, |
| sink.data_ptr<float>(), stream_of(sink)); |
| 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_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]"); |
| chk(blocks >= 1 && blocks <= 1 << 20 && threads >= 32 && threads <= 1024, |
| "bad launch shape"); |
| 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() >= 2, |
| "sink must be CUDA f32 [>=1]"); |
| chk(kind >= 0 && kind <= 6, "kind must be 0..6"); |
| chk(blocks >= 1 && warps >= 1 && warps <= 32, "bad launch shape"); |
| 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_mma_tile(torch::Tensor a, torch::Tensor b, torch::Tensor c) { |
| chk(a.is_cuda() && a.dtype() == torch::kFloat16 && a.is_contiguous() && |
| a.numel() == 256, |
| "a must be contiguous CUDA f16 [16,16]"); |
| chk(b.is_cuda() && b.dtype() == torch::kFloat16 && b.is_contiguous() && |
| b.numel() == 256, |
| "b must be contiguous CUDA f16 [16,16] col-major (pass B.t())"); |
| chk(c.is_cuda() && c.dtype() == torch::kFloat32 && c.is_contiguous() && |
| c.numel() == 256, |
| "c must be contiguous CUDA f32 [16,16]"); |
| mk_mma_tile_launch(a.data_ptr(), b.data_ptr(), c.data_ptr<float>(), |
| stream_of(c)); |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); |
| } |
|
|
| 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(mode >= 0 && mode <= 3, "mode must be 0..3"); |
| 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]"); |
| chk(sink.is_cuda() && sink.dtype() == torch::kFloat32, "sink f32"); |
| 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; |
| } |
|
|
|
|
| void mk_smem_bw(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_smem_bw_launch((int)blocks, (int)threads, iters, |
| sink.data_ptr<float>(), stream_of(sink)); |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); |
| } |
|
|
| void mk_smem_stride(int64_t stride, 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]"); |
| chk(stride >= 1 && stride <= 64, "stride in [1, 64]"); |
| chk(blocks >= 1 && threads >= 32 && threads <= 1024, "bad launch shape"); |
| mk_smem_stride_launch((int)stride, (int)blocks, (int)threads, iters, |
| sink.data_ptr<float>(), stream_of(sink)); |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); |
| } |
|
|
| int64_t mk_mma_fed(torch::Tensor gbuf, int64_t blocks, int64_t warps, |
| int64_t iters, torch::Tensor sink) { |
| chk(gbuf.is_cuda() && gbuf.is_contiguous() && |
| gbuf.dtype() == torch::kFloat16, |
| "gbuf must be contiguous CUDA f16"); |
| chk(gbuf.numel() >= 4096 && gbuf.numel() % 512 == 0, |
| "gbuf numel must be >= 4096 and a multiple of 512"); |
| chk(sink.is_cuda() && sink.dtype() == torch::kFloat32 && sink.numel() >= 2, |
| "sink must be CUDA f32 [>=2]"); |
| chk(blocks >= 1 && warps >= 1 && warps <= 16, "bad launch shape"); |
| cudaDeviceProp prop; |
| cudaGetDeviceProperties(&prop, sink.device().index()); |
| int r = mk_mma_fed_launch(gbuf.data_ptr(), gbuf.numel(), (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; |
| } |
|
|
| double mk_launch_raw_py(int64_t n, int64_t mode, |
| torch::Tensor device_anchor) { |
| chk(n >= 1 && n <= 1000000, "n in [1, 1e6]"); |
| chk(mode == 0 || mode == 1, "mode must be 0|1"); |
| const at::cuda::CUDAGuard guard(device_anchor.device()); |
| double s = 0.0; |
| int r = mk_launch_raw((int)n, (int)mode, &s); |
| return r == 0 ? s : -1.0; |
| } |
|
|
| void mk_stamped_demo(torch::Tensor buf, torch::Tensor out, |
| int64_t fma_iters, torch::Tensor stamps, |
| int64_t blocks, int64_t threads) { |
| chk(buf.is_cuda() && buf.dtype() == torch::kFloat32, "buf f32"); |
| chk(out.numel() >= buf.numel(), "out >= buf"); |
| chk(stamps.is_cuda() && stamps.dtype() == torch::kInt64 && |
| stamps.dim() == 2 && stamps.size(0) >= blocks, |
| "stamps i64 [blocks, slots]"); |
| mk_stamped_demo_launch( |
| buf.const_data_ptr<float>(), out.data_ptr<float>(), buf.numel(), |
| fma_iters, |
| reinterpret_cast<long long*>(stamps.data_ptr<int64_t>()), |
| (int)stamps.size(1), (int)blocks, (int)threads, stream_of(buf)); |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); |
| } |
|
|
| TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { |
| ops.def( |
| "mk_smem_bw(int blocks, int threads, int iters, Tensor! sink) -> ()"); |
| ops.impl("mk_smem_bw", torch::kCUDA, &mk_smem_bw); |
| ops.def( |
| "mk_stamped_demo(Tensor buf, Tensor! out, int fma_iters," |
| " Tensor! stamps, int blocks, int threads) -> ()"); |
| ops.impl("mk_stamped_demo", torch::kCUDA, &mk_stamped_demo); |
| ops.def("mk_triad(Tensor a, Tensor b, Tensor! c, int width, float s) -> ()"); |
| ops.impl("mk_triad", torch::kCUDA, &mk_triad); |
| ops.def( |
| "mk_triad_passes(Tensor a, Tensor b, Tensor! c, float s, int passes)" |
| " -> ()"); |
| ops.impl("mk_triad_passes", torch::kCUDA, &mk_triad_passes); |
| ops.def("mk_read(Tensor a, Tensor! sink) -> ()"); |
| ops.impl("mk_read", torch::kCUDA, &mk_read); |
| ops.def("mk_write(Tensor! c, float v) -> ()"); |
| ops.impl("mk_write", torch::kCUDA, &mk_write); |
| ops.def( |
| "mk_gather(Tensor a, Tensor idx, Tensor! c, int threads," |
| " int dyn_smem) -> ()"); |
| ops.impl("mk_gather", torch::kCUDA, &mk_gather); |
| ops.def( |
| "mk_fma_f64(int blocks, int threads, int iters, Tensor! sink) -> ()"); |
| ops.impl("mk_fma_f64", torch::kCUDA, &mk_fma_f64); |
| ops.def( |
| "mk_fma_dep(int blocks, int threads, int dyn_smem, int iters," |
| " Tensor! sink) -> ()"); |
| ops.impl("mk_fma_dep", torch::kCUDA, &mk_fma_dep); |
| ops.def( |
| "mk_fma_ilp(int chains, int blocks, int threads, int iters," |
| " Tensor! sink) -> ()"); |
| ops.impl("mk_fma_ilp", torch::kCUDA, &mk_fma_ilp); |
| ops.def( |
| "mk_smem_stride(int stride, int blocks, int threads, int iters," |
| " Tensor! sink) -> ()"); |
| ops.impl("mk_smem_stride", torch::kCUDA, &mk_smem_stride); |
| ops.def( |
| "mk_mma_fed(Tensor gbuf, int blocks, int warps, int iters," |
| " Tensor! sink) -> int"); |
| ops.impl("mk_mma_fed", torch::kCUDA, &mk_mma_fed); |
| ops.def("mk_launch_raw(int n, int mode, Tensor device_anchor) -> float"); |
| ops.impl("mk_launch_raw", torch::kCUDA, &mk_launch_raw_py); |
| ops.def("mk_chase_global(Tensor ring, int hops, Tensor! out) -> ()"); |
| ops.impl("mk_chase_global", torch::kCUDA, &mk_chase_global); |
| ops.def("mk_chase_shared(Tensor ring, int hops, Tensor! out) -> ()"); |
| ops.impl("mk_chase_shared", torch::kCUDA, &mk_chase_shared); |
| ops.def( |
| "mk_fma_f32(int blocks, int threads, int iters, Tensor! sink) -> ()"); |
| ops.impl("mk_fma_f32", torch::kCUDA, &mk_fma_f32); |
| ops.def( |
| "mk_mma(int kind, int blocks, int warps, int iters, Tensor! sink)" |
| " -> int"); |
| ops.impl("mk_mma", torch::kCUDA, &mk_mma); |
| ops.def("mk_mma_tile(Tensor a, Tensor b, Tensor! c) -> ()"); |
| ops.impl("mk_mma_tile", torch::kCUDA, &mk_mma_tile); |
| ops.def( |
| "mk_atomics(int mode, int blocks, int threads, int iters, Tensor! buf," |
| " int slots) -> ()"); |
| ops.impl("mk_atomics", torch::kCUDA, &mk_atomics); |
| ops.def( |
| "mk_spin(int blocks, int threads, int ticks, Tensor! out," |
| " Tensor! sink) -> ()"); |
| ops.impl("mk_spin", torch::kCUDA, &mk_spin); |
| ops.def("mk_empty(Tensor device_anchor) -> ()"); |
| ops.impl("mk_empty", torch::kCUDA, &mk_empty); |
| ops.def( |
| "mk_barrier(int blocks, int threads, int iters, Tensor! out) -> int"); |
| ops.impl("mk_barrier", torch::kCUDA, &mk_barrier); |
| ops.def( |
| "mk_occupancy(int threads, int dyn_smem, Tensor device_anchor) -> int"); |
| ops.impl("mk_occupancy", torch::kCUDA, &mk_occupancy); |
| } |
|
|
| REGISTER_EXTENSION(TORCH_EXTENSION_NAME) |
|
|
