Add cute cache

test/test_cute.py

@@ -0,0 +1,153 @@
+#!/usr/bin/env python3
+"""Test script to demonstrate CUTE kernel caching"""
+
+import torch
+from operator import add, mul
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+import pytest
+
+from transformer_nuggets.cute import cute_compile_and_cache, get_cache_stats, clear_cute_cache
+
+
+@cute.kernel
+def simple_add_kernel(
+    gA: cute.Tensor,
+    gB: cute.Tensor,
+    gC: cute.Tensor,
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, _, _ = cute.arch.block_idx()
+    bdim, _, _ = cute.arch.block_dim()
+
+    thread_idx = bidx * bdim + tidx
+
+    m, n = gA.shape
+    ni = thread_idx % n
+    mi = thread_idx // n
+
+    a_val = gA[mi, ni]
+    b_val = gB[mi, ni]
+
+    gC[mi, ni] = a_val + b_val
+
+
+@cute.kernel
+def simple_mul_kernel(
+    gA: cute.Tensor,
+    gB: cute.Tensor,
+    gC: cute.Tensor,
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, _, _ = cute.arch.block_idx()
+    bdim, _, _ = cute.arch.block_dim()
+
+    thread_idx = bidx * bdim + tidx
+
+    m, n = gA.shape
+    ni = thread_idx % n
+    mi = thread_idx // n
+
+    a_val = gA[mi, ni]
+    b_val = gB[mi, ni]
+
+    gC[mi, ni] = a_val * b_val
+
+
+def cached_elementwise(
+    op,
+    mA: cute.Tensor,
+    mB: cute.Tensor,
+    mC: cute.Tensor,
+):
+    # Define the kernel with @cute.jit
+    @cute.jit
+    def add_kernel(mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
+        num_threads_per_block = 256
+        m, n = mA.shape
+
+        simple_add_kernel(mA, mB, mC).launch(
+            grid=((m * n) // num_threads_per_block, 1, 1), block=(num_threads_per_block, 1, 1)
+        )
+
+    @cute.jit
+    def mul_kernel(mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
+        num_threads_per_block = 256
+        m, n = mA.shape
+
+        simple_mul_kernel(mA, mB, mC).launch(
+            grid=((m * n) // num_threads_per_block, 1, 1), block=(num_threads_per_block, 1, 1)
+        )
+
+    # Use explicit caching based on operation
+    if op == add:
+        compiled_kernel = cute_compile_and_cache(add_kernel, mA, mB, mC)
+        return compiled_kernel(mA, mB, mC)
+    elif op == mul:
+        compiled_kernel = cute_compile_and_cache(mul_kernel, mA, mB, mC)
+        return compiled_kernel(mA, mB, mC)
+    else:
+        raise ValueError(f"Unsupported operation: {op}")
+
+
+def test():
+    print("CUTE Kernel Caching Demo")
+    print("=" * 50)
+
+    # Create test tensors
+    M, N = 1024, 1024
+    a = torch.randn(M, N, device="cuda", dtype=torch.float16)
+    b = torch.randn(M, N, device="cuda", dtype=torch.float16)
+    c = torch.zeros(M, N, device="cuda", dtype=torch.float16)
+
+    # Convert to CUTE tensors
+    a_ = from_dlpack(a, assumed_align=16)
+    b_ = from_dlpack(b, assumed_align=16)
+    c_ = from_dlpack(c, assumed_align=16)
+
+    print("\n1. First call with 'add' - should compile:")
+    cached_elementwise(add, a_, b_, c_)
+    torch.testing.assert_close(c, a + b)
+    print("✓ Addition result correct")
+
+    print("\n2. Second call with 'add' - should hit cache:")
+    c.zero_()
+    cached_elementwise(add, a_, b_, c_)
+
+    print("\n3. First call with 'mul' - should compile (different op):")
+    c.zero_()
+    cached_elementwise(mul, a_, b_, c_)
+    torch.testing.assert_close(c, a * b)
+    print("✓ Multiplication result correct")
+
+    print("\n4. Different tensor sizes - should compile:")
+    a2 = torch.randn(512, 512, device="cuda", dtype=torch.float16)
+    b2 = torch.randn(512, 512, device="cuda", dtype=torch.float16)
+    c2 = torch.zeros(512, 512, device="cuda", dtype=torch.float16)
+
+    a2_ = from_dlpack(a2, assumed_align=16)
+    b2_ = from_dlpack(b2, assumed_align=16)
+    c2_ = from_dlpack(c2, assumed_align=16)
+
+    cached_elementwise(add, a2_, b2_, c2_)
+
+    print("\n5. Same size as #4 - should hit cache:")
+    cached_elementwise(add, a2_, b2_, c2_)
+
+    print("\n" + "=" * 50)
+    print("Final Cache Statistics:")
+    stats = get_cache_stats()
+    print(f"  Total calls: {stats['total']}")
+    print(f"  Cache hits: {stats['hits']}")
+    print(f"  Cache misses: {stats['misses']}")
+    print(f"  Hit rate: {stats['hit_rate']:.2%}")
+    print(f"  Unique kernels cached: {stats['cache_size']}")
+
+    print("\nClearing cache...")
+    clear_cute_cache()
+    stats = get_cache_stats()
+    print(f"  Cache size after clear: {stats['cache_size']}")
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

transformer_nuggets/cute/__init__.py

@@ -0,0 +1,6 @@
+from transformer_nuggets.cute.cache import (
+    cute_compile_and_cache,
+    clear_cute_cache,
+    get_cache_stats,
+    set_cache_size,
+)

transformer_nuggets/cute/add.py

@@ -7,6 +7,10 @@
 from cutlass.cute.runtime import from_dlpack
 
 from transformer_nuggets.utils.benchmark import benchmark_cuda_function_in_microseconds
+from transformer_nuggets.cute.cache import cute_compile_and_cache, get_cache_stats
+
+
+# init_logging()
 
 
 @cute.kernel
@@ -31,6 +35,17 @@ def naive_elementwise_add_kernel(
     gC[mi, ni] = a_val + b_val
 
 
+@cute.jit
+def naive_elementwise_add(mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
+    num_threads_per_block = 256
+    m, n = mA.shape
+
+    kernel = naive_elementwise_add_kernel(mA, mB, mC)
+    kernel.launch(
+        grid=((m * n) // num_threads_per_block, 1, 1), block=(num_threads_per_block, 1, 1)
+    )
+
+
 @cute.kernel
 def elementwise_apply_kernel(
     op: cutlass.Constexpr,
@@ -62,23 +77,7 @@ def elementwise_apply_kernel(
 
 
 @cute.jit
-def naive_elementwise_add(mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
-    num_threads_per_block = 256
-    m, n = mA.shape
-
-    kernel = naive_elementwise_add_kernel(mA, mB, mC)
-    kernel.launch(
-        grid=((m * n) // num_threads_per_block, 1, 1), block=(num_threads_per_block, 1, 1)
-    )
-
-
-@cute.jit
-def elementwise_op(
-    op: cutlass.Constexpr,
-    mA: cute.Tensor,
-    mB: cute.Tensor,
-    mC: cute.Tensor,
-):
+def elem_kernel(op: cutlass.Constexpr, mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
     thr_layout = cute.make_layout((4, 32), stride=(32, 1))
     val_layout = cute.make_layout((4, 8), stride=(8, 1))
     tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
@@ -93,13 +92,41 @@ def elementwise_op(
     )
 
 
-def benchmark(callable, *, num_warmups=5, num_iterations=200):
+def elementwise_op(
+    op: cutlass.Constexpr,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    c: torch.Tensor,
+    *,
+    assumed_align: int = 16,
+):
+    """
+    Apply an elementwise operation using CUTE kernels.
+
+    Args:
+        op: Cutlass constexpr operation (e.g., add, mul)
+        a: Input tensor A (PyTorch tensor)
+        b: Input tensor B (PyTorch tensor)
+        c: Output tensor C (PyTorch tensor, will be modified in-place)
+        assumed_align: Memory alignment assumption for dlpack conversion
+    """
+    # Convert PyTorch tensors to CUTE tensors
+    mA = from_dlpack(a, assumed_align=assumed_align)
+    mB = from_dlpack(b, assumed_align=assumed_align)
+    mC = from_dlpack(c, assumed_align=assumed_align)
+
+    # Compile and execute the kernel
+    compiled_kernel = cute_compile_and_cache(elem_kernel, op, mA, mB, mC)
+    return compiled_kernel(mA, mB, mC)
+
+
+def benchmark(callable, tensor_a, *, num_warmups=5, num_iterations=200):
     time = benchmark_cuda_function_in_microseconds(callable)
     avg_time = time / 1e3
 
     print(f"Average execution time: {avg_time:.4f} ms")
 
-    total_bytes = 3 * a.numel() * a.element_size()
+    total_bytes = 3 * tensor_a.numel() * tensor_a.element_size()
     throughput_gb_s = total_bytes / (avg_time / 1000) / 1e9
     print(f"Throughput: {throughput_gb_s:.2f} GB/s")
     print()
@@ -112,28 +139,38 @@ def benchmark(callable, *, num_warmups=5, num_iterations=200):
     b = torch.randn(M, N, device="cuda", dtype=torch.float16)
     c = torch.zeros(M, N, device="cuda", dtype=torch.float16)
 
-    a_ = from_dlpack(a, assumed_align=16)
-    b_ = from_dlpack(b, assumed_align=16)
-    c_ = from_dlpack(c, assumed_align=16)
-
-    elementwise_op(mul, a_, b_, c_)
+    # Test the new API that takes regular PyTorch tensors
+    elementwise_op(mul, a, b, c)
     torch.testing.assert_close(c, mul(a, b))
     print("✓ Multiplication test passed")
 
     c.zero_()
 
+    # For naive kernel, we still need CUTE tensors
+    a_ = from_dlpack(a, assumed_align=16)
+    b_ = from_dlpack(b, assumed_align=16)
+    c_ = from_dlpack(c, assumed_align=16)
     naive_elementwise_add_compiled = cute.compile(naive_elementwise_add, a_, b_, c_)
-    elementwise_add_compiled = cute.compile(elementwise_op, add, a_, b_, c_)
+
+    # Note: elementwise_op is now cached automatically!
+    # First call will compile, subsequent calls use cache
 
     print("\n" + "=" * 50)
     print("BENCHMARK RESULTS")
     print("=" * 50 + "\n")
 
     print("1. Naive elementwise add kernel:")
-    benchmark(partial(naive_elementwise_add_compiled, a_, b_, c_))
+    benchmark(partial(naive_elementwise_add_compiled, a_, b_, c_), a)
 
     print("2. PyTorch add (baseline):")
-    benchmark(partial(torch.add, a, b, out=c))
+    benchmark(partial(torch.add, a, b, out=c), a)
+
+    print("3. Optimized elementwise add kernel (with caching):")
+    benchmark(partial(elementwise_op, add, a, b, c, assumed_align=128), a)
 
-    print("3. Optimized elementwise add kernel:")
-    benchmark(partial(elementwise_add_compiled, a_, b_, c_))
+    print("\nCache Statistics:")
+    stats = get_cache_stats()
+    print(f"  Cache hits: {stats['hits']}")
+    print(f"  Cache misses: {stats['misses']}")
+    print(f"  Hit rate: {stats['hit_rate']:.2%}")
+    print(f"  Cache size: {stats['cache_size']}")