45 lines
1.5 KiB
Python
45 lines
1.5 KiB
Python
"""Absolute minimum: just Int32 constants and Float32 comparisons."""
|
|
import torch
|
|
import cutlass.cute as cute
|
|
import cutlass.torch as cutlass_torch
|
|
from cutlass.cute.typing import Float32, Int32
|
|
import sys
|
|
|
|
test = sys.argv[1] if len(sys.argv) > 1 else "just_store"
|
|
|
|
@cute.kernel
|
|
def just_store(inp: cute.Tensor, out: cute.Tensor):
|
|
tidx, _, _ = cute.arch.thread_idx()
|
|
if tidx == Int32(0):
|
|
cute.arch.store(out.iterator, Int32(42))
|
|
|
|
@cute.kernel
|
|
def load_and_store(inp: cute.Tensor, out: cute.Tensor):
|
|
tidx, _, _ = cute.arch.thread_idx()
|
|
if tidx == Int32(0):
|
|
x = cute.arch.load(inp.iterator, Float32)
|
|
cute.arch.store(out.iterator, Int32(1))
|
|
|
|
@cute.kernel
|
|
def float_cmp(inp: cute.Tensor, out: cute.Tensor):
|
|
tidx, _, _ = cute.arch.thread_idx()
|
|
if tidx == Int32(0):
|
|
x = cute.arch.load(inp.iterator, Float32)
|
|
r = Int32(0)
|
|
if x > Float32(0.5):
|
|
r = Int32(1)
|
|
cute.arch.store(out.iterator, r)
|
|
|
|
KERNELS = {"just_store": just_store, "load_and_store": load_and_store, "float_cmp": float_cmp}
|
|
k = KERNELS[test]
|
|
|
|
if __name__ == "__main__":
|
|
x = torch.tensor([3.7], dtype=torch.float32, device='cuda')
|
|
o = torch.zeros(1, dtype=torch.int32, device='cuda')
|
|
xc = cutlass_torch.from_dlpack(x).mark_layout_dynamic(leading_dim=0)
|
|
oc = cutlass_torch.from_dlpack(o).mark_layout_dynamic(leading_dim=0)
|
|
print(f"Test: {test}")
|
|
compiled = cute.compile(k, xc, oc)
|
|
compiled(xc, oc)
|
|
print(f"Result: {o.item()}")
|