resnet infer

Author: qinyiqun

avg_pool2d.py

@@ -0,0 +1,101 @@
+from contextlib import contextmanager
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import triton
+
+import ops.ninetoothed.torch
+import ops.triton.torch
+import ntops.torch
+
+class AvgPool2d(nn.Module):
+    avg_pool2d = None
+
+    def __init__(self, other):
+        super().__init__()
+
+        self.__dict__ = other.__dict__
+
+    def forward(self, input):
+        def _pair(x):
+            return (x, x) if isinstance(x, int) else x
+        
+        return type(self).avg_pool2d(
+            input,
+            kernel_size=_pair(self.kernel_size),
+            stride=_pair(self.stride) if self.stride else _pair(self.kernel_size),
+            padding=_pair(self.padding),
+            ceil_mode=self.ceil_mode,
+            count_include_pad=self.count_include_pad,
+            divisor_override=self.divisor_override,
+        )
+
+
+@contextmanager
+def avg_pool2d_backend(backend_name):
+    _prev_impl = AvgPool2d.avg_pool2d
+
+    if backend_name == "ninetoothed":
+        impl = ntops.torch.avg_pool2d
+    elif backend_name == "triton":
+        impl = ops.triton.torch.avg_pool2d
+    elif backend_name == "torch":
+        impl = F.avg_pool2d
+    else:
+        raise ValueError(f"unknown backend: `{backend_name}`")
+
+    AvgPool2d.avg_pool2d = impl
+
+    try:
+        yield
+    finally:
+        AvgPool2d.avg_pool2d = _prev_impl
+
+if __name__ == "__main__":
+    torch.manual_seed(0)
+
+    input_shape = (32, 3, 64, 64)
+    window_shape = (3, 3)
+
+    input = torch.randn(input_shape, dtype=torch.float16, device="cuda")
+
+    ninetoothed_output = ops.ninetoothed.torch.avg_pool2d(input, window_shape)
+    torch_output = F.avg_pool2d(input, window_shape, ceil_mode=True)
+
+    print(ninetoothed_output)
+    print(torch_output)
+
+    if torch.allclose(ninetoothed_output, torch_output):
+        print("✅ NineToothed and PyTorch match.")
+    else:
+        print("❌ NineToothed and PyTorch differ.")
+
+    @triton.testing.perf_report(
+        triton.testing.Benchmark(
+            x_names=["h", "w"],
+            x_vals=[8 * i for i in range(2, 33)],
+            line_arg="provider",
+            line_vals=["ninetoothed", "torch"],
+            line_names=["NineToothed", "PyTorch"],
+            styles=[("blue", "-"), ("green", "-")],
+            ylabel="ms",
+            plot_name="avg-pool2d-performance",
+            args={},
+        )
+    )
+    def benchmark(h, w, provider):
+        n, c, h, w = 64, 64, h, w
+        r, s = 3, 3
+        dtype = torch.float16
+        device = "cuda"
+        input = torch.randn((n, c, h, w), dtype=dtype, device=device)
+        window_shape = (r, s)
+
+        if provider == "ninetoothed":
+            ms = triton.testing.do_bench(lambda: ops.ninetoothed.torch.avg_pool2d(input, window_shape))
+        elif provider == "torch":
+            ms = triton.testing.do_bench(lambda: F.avg_pool2d(input, window_shape))
+
+        return ms
+
+    benchmark.run(show_plots=True, print_data=True, save_path=".")

conv2d.py

@@ -1,75 +1,84 @@
+from contextlib import contextmanager
 import torch
 import torch.nn.functional as F
 import triton
+import torch.nn as nn
 
 import ops.ninetoothed.torch
 import ops.triton.torch
+import ntops.torch
+
+class Conv2d(nn.Module):
+    conv2d = None
+
+    def __init__(self, other):
+        super().__init__()
+
+        self.__dict__ = other.__dict__
+
+    def forward(self, input):
+        def _pair(x):
+            return x if isinstance(x, (tuple, list)) else (x, x)
+        
+        stride = _pair(self.stride)
+        padding = _pair(self.padding)
+        dilation = _pair(self.dilation)
+        
+        return type(self).conv2d(
+            input,
+            self.weight,
+            self.bias,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=self.groups,
+        )
+    
+@contextmanager
+def conv2d_backend(backend_name):
+    _prev_impl = Conv2d.conv2d
+    if backend_name == "ninetoothed":
+        impl = ntops.torch.conv2d
+    elif backend_name == "torch":
+        impl = F.conv2d
+    else:
+        raise ValueError(f"unknown backend: `{backend_name}`")
+
+    Conv2d.conv2d = impl
+
+    try:
+        yield
+    finally:
+        Conv2d.conv2d = _prev_impl
 
 if __name__ == "__main__":
     torch.manual_seed(0)
 
-    n, c, h, w = 4, 3, 224, 224
-    k, _, r, s = 8, c, 3, 3
+    n, c, h, w = 2, 3, 112, 112
+    k, _, r, s = 4, c, 3, 3
     dtype = torch.float16
     device = "cuda"
 
-    input = torch.randn(n, c, h, w, dtype=dtype, device=device)
-    filter = torch.randn(k, c, r, s, dtype=dtype, device=device)
-
-    ninetoothed_output = ops.ninetoothed.torch.conv2d(input, filter)
-    torch_output = F.conv2d(input, filter)
-    triton_output = ops.triton.torch.conv2d(input, filter)
+    input = torch.randn((n, c, h, w), dtype=dtype, device=device)
+    weight = torch.randn((k, c, r, s), dtype=dtype, device=device)
+    bias = torch.randn((k,), dtype=dtype, device=device)
+    stride = 3
+    dilation = 1
+    
+    ninetoothed_output = ops.ninetoothed.torch.conv2d(
+        input, weight, bias=bias, stride=stride, dilation=dilation
+    )
+    reference_output = F.conv2d(
+        input, weight, bias=bias, stride=stride, dilation=dilation
+    )
 
+    
     print(ninetoothed_output)
-    print(torch_output)
-    print(triton_output)
+    print(reference_output)
+    
 
-    if torch.allclose(ninetoothed_output, torch_output, atol=0.01, rtol=0.01):
+    if torch.allclose(ninetoothed_output, reference_output, atol=0.01, rtol=0.01):
         print("✅ NineToothed and PyTorch match.")
     else:
         print("❌ NineToothed and PyTorch differ.")
-    if torch.allclose(ninetoothed_output, triton_output, atol=0, rtol=0):
-        print("✅ NineToothed and Triton match.")
-    else:
-        print("❌ NineToothed and Triton differ.")
-
-    @triton.testing.perf_report(
-        triton.testing.Benchmark(
-            x_names=["n"],
-            x_vals=[2**i for i in range(1, 11)],
-            x_log=True,
-            line_arg="provider",
-            line_vals=["ninetoothed", "torch", "triton"],
-            line_names=["NineToothed", "PyTorch", "Triton"],
-            styles=[("blue", "-"), ("green", "-"), ("orange", "-")],
-            ylabel="ms",
-            plot_name="conv2d-performance",
-            args={},
-        )
-    )
-    def benchmark(n, provider):
-        _, c, h, w = n, 512, 14, 14
-        k, _, r, s = 512, c, 3, 3
-
-        input = torch.randn((n, c, h, w), dtype=dtype, device=device)
-        filter = torch.randn((k, c, r, s), dtype=dtype, device=device)
-
-        ninetoothed_output = ops.ninetoothed.torch.conv2d(input, filter)
-        torch_output = F.conv2d(input, filter)
-        triton_output = ops.triton.torch.conv2d(input, filter)
-
-        assert torch.allclose(ninetoothed_output, torch_output, atol=0.01, rtol=0.01)
-        assert torch.allclose(ninetoothed_output, triton_output, atol=0, rtol=0)
-
-        if provider == "ninetoothed":
-            ms = triton.testing.do_bench(
-                lambda: ops.ninetoothed.torch.conv2d(input, filter)
-            )
-        elif provider == "torch":
-            ms = triton.testing.do_bench(lambda: F.conv2d(input, filter))
-        elif provider == "triton":
-            ms = triton.testing.do_bench(lambda: ops.triton.torch.conv2d(input, filter))
-
-        return ms
-
-    benchmark.run(show_plots=True, print_data=True, save_path=".")
+    

max_pool2d.py

@@ -1,54 +1,79 @@
-import math
-
-import ninetoothed
-import ninetoothed.language as ntl
+from contextlib import contextmanager
 import torch
+import torch.nn as nn
 import torch.nn.functional as F
 import triton
-from ninetoothed import Symbol, Tensor
-
-
-def arrangement(input, output):
-    BLOCK_SIZE = Symbol("BLOCK_SIZE", meta=True)
-
-    WINDOW_HEIGHT = Symbol("WINDOW_HEIGHT", constexpr=True, upper_bound=16)
-    WINDOW_WIDTH = Symbol("WINDOW_WIDTH", constexpr=True, upper_bound=16)
-
-    input_arranged = input.tile((1, 1, WINDOW_HEIGHT, WINDOW_WIDTH))
-    input_arranged = input_arranged.ravel()
-    input_arranged = input_arranged.flatten(end_dim=4).flatten(start_dim=1)
-    input_arranged = input_arranged.tile((BLOCK_SIZE, -1))
-
-    output_arranged = output.tile((1, 1, 1, 1))
-    output_arranged = output_arranged.ravel()
-    output_arranged = output_arranged.flatten(end_dim=4).flatten(start_dim=1)
-    output_arranged = output_arranged.tile((BLOCK_SIZE, -1))
-    output_arranged.dtype = output_arranged.dtype.squeeze(1)
-
-    return input_arranged, output_arranged
-
-
-def application(input, output):
-    output = ntl.max(input, axis=1)  # noqa: F841
-
-
-max_pool2d_kernel = ninetoothed.make(
-    arrangement, application, (Tensor(4, other=float("-inf")), Tensor(4))
-)
 
+import ops.ninetoothed.torch
+import ops.triton.torch
+
+import ntops.torch
+
+class MaxPool2d(nn.Module):
+    max_pool2d = None
+
+    def __init__(self, other):
+        super().__init__()
+
+        self.__dict__ = other.__dict__
+
+    # def forward(self, input):
+    #     print("using ninetoothed max_pool2d kernel")
+    #     return type(self).max_pool2d(input)
+    def forward(self, input):
+        
+        # 处理 kernel_size（确保是 tuple）
+        kernel_size = self.kernel_size
+        if isinstance(kernel_size, int):
+            kernel_size = (kernel_size, kernel_size)
+        
+        # 处理 stride（None 表示默认等于 kernel_size）
+        stride = self.stride
+        if stride is None:
+            stride = kernel_size
+        elif isinstance(stride, int):
+            stride = (stride, stride)
+        
+        # 处理 padding
+        padding = self.padding
+        if isinstance(padding, int):
+            padding = (padding, padding)
+        
+        # 处理 dilation
+        dilation = self.dilation
+        if isinstance(dilation, int):
+            dilation = (dilation, dilation)
+        
+        return type(self).max_pool2d(
+            input,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            ceil_mode=self.ceil_mode,
+            return_indices=self.return_indices,
+        )
 
-def max_pool2d(input, window_shape):
-    n, c, h, w = input.shape
-    r, s = window_shape
-    p = math.ceil((h - r) / r + 1)
-    q = math.ceil((w - s) / s + 1)
 
-    output = torch.empty(n, c, p, q, dtype=input.dtype, device=input.device)
+@contextmanager
+def max_pool2d_backend(backend_name):
+    _prev_impl = MaxPool2d.max_pool2d
 
-    max_pool2d_kernel(input, output, WINDOW_HEIGHT=r, WINDOW_WIDTH=s)
+    if backend_name == "ninetoothed":
+        impl = ntops.torch.max_pool2d
+    elif backend_name == "triton":
+        impl = ops.triton.torch.max_pool2d
+    elif backend_name == "torch":
+        impl = F.max_pool2d
+    else:
+        raise ValueError(f"unknown backend: `{backend_name}`")
 
-    return output
+    MaxPool2d.max_pool2d = impl
 
+    try:
+        yield
+    finally:
+        MaxPool2d.max_pool2d = _prev_impl
 
 if __name__ == "__main__":
     torch.manual_seed(0)
@@ -58,7 +83,7 @@ def max_pool2d(input, window_shape):
 
     input = torch.randn(input_shape, dtype=torch.float16, device="cuda")
 
-    ninetoothed_output = max_pool2d(input, window_shape)
+    ninetoothed_output = ops.ninetoothed.torch.max_pool2d(input, window_shape)
     torch_output = F.max_pool2d(input, window_shape, ceil_mode=True)
 
     print(ninetoothed_output)
@@ -91,7 +116,7 @@ def benchmark(h, w, provider):
         window_shape = (r, s)
 
         if provider == "ninetoothed":
-            ms = triton.testing.do_bench(lambda: max_pool2d(input, window_shape))
+            ms = triton.testing.do_bench(lambda: ops.ninetoothed.torch.max_pool2d(input, window_shape))
         elif provider == "torch":
             ms = triton.testing.do_bench(lambda: F.max_pool2d(input, window_shape))