diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index bd2373e3dd8..ca8f9eee753 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -176,6 +176,7 @@ fbcode_target( "triton/kernels/sdpa.py", "triton/kernels/topk.py", "triton/kernels/tq4_sdpa.py", + "triton/kernels/triton_mutation_cache.py", ], visibility = [ "//executorch/backends/cuda/...", diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index f28faabb624..aad054fa288 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -42,6 +42,7 @@ ReplaceInt64FloorDivWithFloatPass, ) from executorch.backends.cuda.target_smem import target_smem_context +from executorch.backends.cuda.triton.kernels import triton_mutation_cache from executorch.backends.cuda.triton.replacement_pass import ( ReplaceEdgeOpWithTritonOpPass, ) @@ -958,6 +959,7 @@ def _combined(): # them. SDPA ops already replaced by Triton kernels via # `ReplaceEdgeOpWithTritonOpPass` are unaffected. stack.enter_context(torch.nn.attention.sdpa_kernel([SDPBackend.MATH])) + stack.enter_context(triton_mutation_cache.enabled()) stack.enter_context(target_smem_context()) # On ROCm, is_big_gpu gates on architecture, not SM count. if torch.version.hip is None: diff --git a/backends/cuda/tests/targets.bzl b/backends/cuda/tests/targets.bzl index 8630a2e2c31..576943f147b 100644 --- a/backends/cuda/tests/targets.bzl +++ b/backends/cuda/tests/targets.bzl @@ -306,6 +306,25 @@ def define_common_targets(is_fbcode = False): ], ) + python_unittest_remote_gpu( + name = "test_triton_mutation_cache", + srcs = [ + "test_triton_mutation_cache.py", + ], + visibility = [ + "//executorch/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/cuda:triton_kernels", + ], + keep_gpu_sections = True, + remote_execution = re_test_utils.remote_execution( + platform = "gpu-remote-execution", + subplatform = "A100-exclusive", + ), + ) + python_unittest_remote_gpu( name = "test_cuda_weight_reuse", srcs = [ diff --git a/backends/cuda/tests/test_triton_mutation_cache.py b/backends/cuda/tests/test_triton_mutation_cache.py new file mode 100644 index 00000000000..61fb675e3b5 --- /dev/null +++ b/backends/cuda/tests/test_triton_mutation_cache.py @@ -0,0 +1,328 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""triton_mutation_cache: repeats of PyTorch's Triton access analysis are reused.""" + +import re +import threading +import time +import unittest +from unittest import mock + +import torch +import triton +import triton.language as tl +from executorch.backends.cuda.triton.kernels import triton_mutation_cache +from torch._higher_order_ops import triton_kernel_wrap +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.library import wrap_triton + + +@triton.jit +def _scale_kernel( + x_ptr, y_ptr, z_ptr, n, scale, WRITE_Z: tl.constexpr, BLOCK: tl.constexpr +): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < n + x = tl.load(x_ptr + offsets, mask=mask) + tl.store(y_ptr + offsets, x * scale, mask=mask) + if WRITE_Z: + tl.store(z_ptr + offsets, x + scale, mask=mask) + + +@triton.jit +def _store_z_kernel( + x_ptr, y_ptr, z_ptr, n, scale, WRITE_Z: tl.constexpr, BLOCK: tl.constexpr +): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < n + x = tl.load(x_ptr + offsets, mask=mask) + tl.store(z_ptr + offsets, x * scale, mask=mask) + + +def _scale(x: torch.Tensor, scale: float) -> torch.Tensor: + y = torch.empty_like(x) + z = torch.empty_like(x) + n = x.numel() + wrap_triton(_scale_kernel)[(triton.cdiv(n, 256),)]( + x, y, z, n, scale, WRITE_Z=False, BLOCK=256 + ) + return y + + +class _Layers(torch.nn.Module): + """Calls the kernel directly, so export traces (and analyzes) each call.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Like a model's layers: the same kernel, dtypes and scalars each time. + for _ in range(4): + x = _scale(x, 0.5) + return x + + +def _written(analysis) -> set: + return {dep.name for dep in analysis.read_writes.writes} + + +def _scale_with_num_warps(x: torch.Tensor, scale: float) -> torch.Tensor: + y = torch.empty_like(x) + z = torch.empty_like(x) + n = x.numel() + wrap_triton(_scale_kernel)[(triton.cdiv(n, 256),)]( + x, y, z, n, scale, WRITE_Z=False, BLOCK=256, num_warps=4 + ) + return y + + +class _NumWarpsLayers(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + for _ in range(4): + x = _scale_with_num_warps(x, 0.5) + return x + + +class TritonMutationCacheTest(unittest.TestCase): + def setUp(self) -> None: + if not torch.cuda.is_available(): + self.skipTest("CUDA required") + self.original = triton_kernel_wrap.identify_accessed_tensors + + def _kwargs(self, mode, **overrides): + with mode: + x = torch.empty(1000, device="cuda") + kwargs = { + "x_ptr": x, + "y_ptr": torch.empty_like(x), + "z_ptr": torch.empty_like(x), + "n": 1000, + "scale": 2.0, + "WRITE_Z": False, + "BLOCK": 256, + } + kwargs.update(overrides) + return kwargs + + def _counting_stand_in(self, delay: float = 0.0): + """A stand-in analysis that counts its calls, to patch in for PyTorch's.""" + calls = [] + + def analysis(kernel, kwargs, tma_descriptor_metadata): + calls.append(1) + time.sleep(delay) + return {"writes": sorted(kwargs)} + + return ( + mock.patch.object( + triton_kernel_wrap, "identify_accessed_tensors", analysis + ), + calls, + ) + + def test_the_context_patches_the_function_pytorch_calls_and_restores_it( + self, + ) -> None: + with triton_mutation_cache.enabled(): + self.assertIsNot( + triton_kernel_wrap.identify_accessed_tensors, self.original + ) + self.assertIs(triton_kernel_wrap.identify_accessed_tensors, self.original) + + def test_first_call_goes_through_pytorch_and_repeats_reuse_it(self) -> None: + mode = FakeTensorMode() + with triton_mutation_cache.enabled() as counts: + analyze = triton_kernel_wrap.identify_accessed_tensors + first = analyze(_scale_kernel, self._kwargs(mode), {}) + # Other fake tensors of the same dtypes and the same scalars. + for _ in range(3): + again = analyze(_scale_kernel, self._kwargs(FakeTensorMode()), {}) + self.assertEqual(again, first) + self.assertIsNot(again, first) + self.assertEqual(counts, {"hits": 3, "misses": 1, "uncached": 0}) + self.assertEqual(first, self.original(_scale_kernel, self._kwargs(mode), {})) + + def test_what_changes_the_kernel_is_not_reused(self) -> None: + mode = FakeTensorMode() + variants = [ + {}, + {"WRITE_Z": True}, + {"BLOCK": 128}, + {"n": 1}, + {"scale": 3.0}, + ] + with mode: + half = torch.empty(1000, device="cuda", dtype=torch.float16) + variants.append({"x_ptr": half}) + with triton_mutation_cache.enabled() as counts: + analyze = triton_kernel_wrap.identify_accessed_tensors + for overrides in variants: + kwargs = self._kwargs(mode, **overrides) + self.assertEqual( + analyze(_scale_kernel, kwargs, {}), + self.original(_scale_kernel, kwargs, {}), + overrides, + ) + self.assertEqual(counts["misses"], len(variants)) + self.assertEqual( + _written( + self.original(_scale_kernel, self._kwargs(mode, WRITE_Z=True), {}) + ), + {"y_ptr", "z_ptr"}, + ) + + def test_kernels_with_the_same_signature_are_kept_apart(self) -> None: + mode = FakeTensorMode() + kwargs = self._kwargs(mode, WRITE_Z=True) + with triton_mutation_cache.enabled(): + analyze = triton_kernel_wrap.identify_accessed_tensors + scale = analyze(_scale_kernel, kwargs, {}) + store = analyze(_store_z_kernel, kwargs, {}) + self.assertEqual(_written(scale), {"y_ptr", "z_ptr"}) + self.assertEqual(_written(store), {"z_ptr"}) + + def test_an_autotuners_first_config_is_part_of_the_key(self) -> None: + tuned = triton.autotune( + configs=[triton.Config({"WRITE_Z": False, "BLOCK": 256})], key=["n"] + )(_scale_kernel) + kwargs = self._kwargs(FakeTensorMode()) + with triton_mutation_cache.enabled(): + analyze = triton_kernel_wrap.identify_accessed_tensors + before = analyze(tuned, kwargs, {}) + tuned.configs = [triton.Config({"WRITE_Z": True, "BLOCK": 256})] + after = analyze(tuned, kwargs, {}) + self.assertEqual(_written(before), {"y_ptr"}) + self.assertEqual(_written(after), {"y_ptr", "z_ptr"}) + + def test_fresh_autotuners_of_one_kernel_share_an_entry(self) -> None: + mode = FakeTensorMode() + with triton_mutation_cache.enabled() as counts: + analyze = triton_kernel_wrap.identify_accessed_tensors + for _ in range(2): + tuned = triton.autotune( + configs=[ + triton.Config({"WRITE_Z": False, "BLOCK": 256}, num_warps=4) + ], + key=["n"], + )(_scale_kernel) + analyze(tuned, self._kwargs(mode), {}) + self.assertEqual(counts["hits"], 1) + self.assertEqual(counts["misses"], 1) + + def test_tma_descriptors_and_unhashable_values_always_go_through_pytorch( + self, + ) -> None: + patch, calls = self._counting_stand_in() + with patch, triton_mutation_cache.enabled() as counts: + analyze = triton_kernel_wrap.identify_accessed_tensors + for _ in range(2): + analyze(_scale_kernel, {"x_ptr": 1}, {"x_ptr": ("stable", ([16],))}) + analyze(_scale_kernel, {"x_ptr": [1, 2]}, {}) + self.assertEqual(len(calls), 4) + self.assertEqual(counts, {"hits": 0, "misses": 0, "uncached": 4}) + + def test_concurrent_calls_of_a_key_analyze_once(self) -> None: + patch, calls = self._counting_stand_in(delay=0.2) + results = [] + with patch, triton_mutation_cache.enabled() as counts: + analyze = triton_kernel_wrap.identify_accessed_tensors + threads = [ + threading.Thread( + target=lambda: results.append(analyze(_scale_kernel, {"n": 7}, {})) + ) + for _ in range(4) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + self.assertEqual(len(calls), 1) + self.assertEqual(results, [{"writes": ["n"]}] * 4) + self.assertEqual(counts, {"hits": 3, "misses": 1, "uncached": 0}) + + def test_real_tensors_are_keyed_by_dtype_and_pointer_alignment(self) -> None: + patch, calls = self._counting_stand_in() + base = torch.empty(64, dtype=torch.float32, device="cuda") + with patch, triton_mutation_cache.enabled() as counts: + analyze = triton_kernel_wrap.identify_accessed_tensors + # base[4:] starts 16 bytes in, so its pointer has base's alignment; + # base[1:] starts 4 bytes in and does not. + for x in (base, base[4:], base[1:], base[1:]): + analyze(_scale_kernel, {"x_ptr": x}, {}) + self.assertEqual(len(calls), 2) + self.assertEqual(counts, {"hits": 2, "misses": 2, "uncached": 0}) + + def test_entries_do_not_outlive_the_context(self) -> None: + patch, calls = self._counting_stand_in() + with patch: + for _ in range(2): + with triton_mutation_cache.enabled(): + analyze = triton_kernel_wrap.identify_accessed_tensors + analyze(_scale_kernel, {"n": 1}, {}) + analyze(_scale_kernel, {"n": 1}, {}) + self.assertEqual(len(calls), 2) + + def test_tensor_shapes_and_strides_do_not_change_the_analysis(self) -> None: + # The key leaves them out; PyTorch's own analysis must not depend on them. + mode = FakeTensorMode() + with mode: + tensors = [ + torch.empty(1000, device="cuda"), + torch.empty(64, 32, device="cuda").t(), + torch.empty(2048, device="cuda")[1:], + ] + first = self.original(_scale_kernel, self._kwargs(mode), {}) + with triton_mutation_cache.enabled(): + analyze = triton_kernel_wrap.identify_accessed_tensors + for x in tensors: + kwargs = self._kwargs(mode, x_ptr=x, y_ptr=x, z_ptr=x) + self.assertEqual(self.original(_scale_kernel, kwargs, {}), first) + self.assertEqual(analyze(_scale_kernel, kwargs, {}), first) + + def test_export_traces_the_same_graph_and_analyzes_once(self) -> None: + x = torch.randn(1000, device="cuda") + + def exported_code(): + program = torch.export.export(_Layers(), (x,)).run_decompositions({}) + # Side-table indices grow with every export; the rest must match. + return re.sub( + r"(kernel_idx|constant_args_idx) = \d+", + r"\1 = N", + program.graph_module.code, + ) + + with triton_mutation_cache.enabled() as counts: + cached = exported_code() + self.assertEqual(counts["misses"], 1) + self.assertGreaterEqual(counts["hits"], 3) + self.assertEqual(cached, exported_code()) + + def test_a_launch_with_num_warps_is_analyzed_once(self) -> None: + x = torch.randn(1000, device="cuda") + with triton_mutation_cache.enabled() as counts: + torch.export.export(_NumWarpsLayers(), (x,)).run_decompositions({}) + self.assertEqual(counts["misses"], 1) + self.assertGreaterEqual(counts["hits"], 3) + + def test_inductor_compile_gives_eager_results_and_reaches_the_analysis( + self, + ) -> None: + x = torch.randn(1000, device="cuda") + with triton_mutation_cache.enabled() as counts: + out = torch.compile(_Layers(), fullgraph=True)(x) + torch.testing.assert_close(out, _Layers()(x)) + self.assertGreaterEqual(counts["misses"], 1) + self.assertEqual(counts["uncached"], 0) + + def test_a_dynamic_shape_export_analyzes_the_symbolic_size_once(self) -> None: + x = torch.randn(1000, device="cuda") + n = torch.export.Dim("n", min=2, max=1 << 20) + with triton_mutation_cache.enabled() as counts: + program = torch.export.export( + _Layers(), (x,), dynamic_shapes=({0: n},) + ).run_decompositions({}) + self.assertEqual(counts["misses"], 1) + self.assertGreaterEqual(counts["hits"], 3) + other = torch.randn(4096, device="cuda") + torch.testing.assert_close(program.module()(other), _Layers()(other)) diff --git a/backends/cuda/triton/kernels/triton_mutation_cache.py b/backends/cuda/triton/kernels/triton_mutation_cache.py new file mode 100644 index 00000000000..096d2b3ee26 --- /dev/null +++ b/backends/cuda/triton/kernels/triton_mutation_cache.py @@ -0,0 +1,131 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Remember PyTorch's analysis of which tensors a user Triton kernel accesses. + +Tracing a call of a user-defined Triton kernel (functionalization in +``torch.export``, and Inductor's mutation propagation and lowering) runs +``torch._higher_order_ops.triton_kernel_wrap.identify_accessed_tensors``: it +generates the kernel's TTIR and walks it to find the tensors the kernel reads +and writes. Nothing caches it, so a model whose every layer launches the same +kernels repeats it for every call site of every exported method. + +The analysis depends only on the traced kernel and on what ``generate_ttir`` +turns the arguments into: a fake or Inductor tensor becomes a fresh tensor of +its dtype, a symbolic value becomes 2, and every other value is used as is. +``enabled()`` caches it for the duration of a compile. Calls this cannot key +exactly (TMA descriptors, unhashable values) always go through PyTorch. +""" + +from __future__ import annotations + +import contextlib +import copy +import threading +from typing import Any, Dict, Iterator, Optional, Tuple + +import torch +import triton + +_UNKEYABLE = object() + +_lazy_types: Optional[Tuple[Tuple[type, ...], Tuple[type, ...]]] = None + + +def _symbolic_and_traced_tensor_types() -> Tuple[Tuple[type, ...], Tuple[type, ...]]: + """The symbolic-value and fake/Inductor tensor types, imported on first use + so that tracing does not load Inductor before it is needed.""" + global _lazy_types + if _lazy_types is None: + import sympy + from torch._inductor.ir import TensorBox + from torch._subclasses.fake_tensor import FakeTensor + + _lazy_types = ( + (torch.SymInt, torch.SymFloat, torch.SymBool, sympy.Expr), + (FakeTensor, TensorBox), + ) + return _lazy_types + + +def _argument_key(value: Any) -> Any: + """What ``generate_ttir`` makes of ``value``, as a hashable key.""" + symbolic, traced_tensor = _symbolic_and_traced_tensor_types() + if isinstance(value, symbolic): + return ("symbolic",) + if isinstance(value, traced_tensor): + return ("tensor", value.dtype) + if isinstance(value, torch.Tensor): + # Triton specializes on the pointer's alignment, so that is what matters. + return ("real tensor", value.dtype, value.data_ptr() % 16) + try: + hash(value) + except TypeError: + return _UNKEYABLE + return (type(value).__module__, type(value).__qualname__, value) + + +def _traced_kernel_and_kwargs(kernel: Any, kwargs: Dict[str, Any]) -> Tuple[Any, Dict]: + """The JIT function ``generate_ttir`` traces, and the arguments it traces it + with: an autotuner's first config is merged into the arguments.""" + if isinstance(kernel, triton.runtime.Autotuner) and kernel.configs: + return kernel.fn, {**kwargs, **kernel.configs[0].kwargs} + return kernel, kwargs + + +def _call_key( + kernel: Any, kwargs: Dict[str, Any], tma_descriptor_metadata +) -> Optional[Tuple[Tuple, Any]]: + """The cache key of a call and the traced kernel it belongs to, or None.""" + if tma_descriptor_metadata: + return None + traced, merged = _traced_kernel_and_kwargs(kernel, kwargs) + arguments = [] + for name in sorted(merged): + key = _argument_key(merged[name]) + if key is _UNKEYABLE: + return None + arguments.append((name, key)) + return (id(traced), tuple(arguments)), traced + + +@contextlib.contextmanager +def enabled() -> Iterator[Dict[str, int]]: + """Caches PyTorch's access analysis while the context is active. Yields the + hit, miss and uncached call counts. Contexts must not overlap: each one + patches the function PyTorch calls.""" + from torch._higher_order_ops import triton_kernel_wrap + + original = triton_kernel_wrap.identify_accessed_tensors + # key -> (traced kernel, result). The kernel is held so its id stays unique. + entries: Dict[Tuple, Tuple[Any, Any]] = {} + counts = {"hits": 0, "misses": 0, "uncached": 0} + lock = threading.RLock() + + def cached(kernel, kwargs, tma_descriptor_metadata): + keyed = _call_key(kernel, kwargs, tma_descriptor_metadata) + if keyed is None: + with lock: + counts["uncached"] += 1 + return original(kernel, kwargs, tma_descriptor_metadata) + key, traced = keyed + with lock: + entry = entries.get(key) + if entry is not None: + counts["hits"] += 1 + return copy.deepcopy(entry[1]) + counts["misses"] += 1 + result = original(kernel, kwargs, tma_descriptor_metadata) + entries[key] = (traced, copy.deepcopy(result)) + return result + + # Both callers (get_mutated_tensors and Inductor's UserDefinedTritonKernel) + # look the function up on this module. + triton_kernel_wrap.identify_accessed_tensors = cached + try: + yield counts + finally: + triton_kernel_wrap.identify_accessed_tensors = original