[general]
name = "megablocks"
universal = false

[torch]
src = [
  "torch-ext/torch_binding.cpp",
  "torch-ext/torch_binding.h"
]

[kernel.megablocks]
backend = "cuda"
cuda-capabilities = [
    "7.0",
    "7.2",
    "7.5",
    "8.0",
    "8.6",
    "8.7",
    "8.9",
    "9.0",
    "10.0",
    "10.1",
    "12.0",
]
depends = ["torch", "cutlass_3_8"]
src = [
    "csrc/new_cumsum.h",
    "csrc/new_cumsum.cu",
    "csrc/new_histogram.h",
    "csrc/new_histogram.cu",
    "csrc/new_indices.h",
    "csrc/new_indices.cu",
    "csrc/new_replicate.cu",
    "csrc/new_replicate.h",
    "csrc/new_sort.h",
    "csrc/new_sort.cu",
]