Revert "Re-land precompile triton templates (#124030)"

pytorchmergebot · pytorchmergebot · commit 2b82345e487a · 2024-04-18T07:21:41.000Z
This reverts commit 030bb13. Reverted #124030 on behalf of https://github.com/DanilBaibak due to Broken trunk ([comment](#124030 (comment)))
diff --git a/test/inductor/test_max_autotune.py b/test/inductor/test_max_autotune.py
@@ -328,8 +328,7 @@ def no_lookup(
             inputs: str,
             benchmark: Callable[[Any], Dict[ChoiceCaller, float]],
         ) -> Dict[ChoiceCaller, float]:
-            if benchmark is not None:
-                return benchmark(choices)
+            return benchmark(choices)
 
         asc = AlgorithmSelectorCache()
 
@@ -427,25 +426,6 @@ def foo(mod, x):
             FileCheck().check_not("extern_kernels.convolution").run(code[0])
             self.assertEqual(conv1x1(input_tensor), out, atol=1e-2, rtol=0)
 
-    @skipIfRocm
-    def test_filled_cache_precompile(self):
-        def fn(a, b, c):
-            a = (a @ b) @ c
-            a, b, c = (t.to(torch.float16) for t in [a, b, c])
-            return (a @ b) @ c
-
-        fn_c = torch.compile(mode="max-autotune-no-cudagraphs")(fn)
-        inputs = [torch.rand([256, 256], device="cuda") for _ in range(3)]
-        from torch._dynamo.utils import counters
-
-        self.assertEqual(fn(*inputs), fn_c(*inputs), atol=1e-2, rtol=1e-2)
-
-        torch._dynamo.reset()
-        counters.clear()
-
-        fn_c = torch.compile(mode="max-autotune-no-cudagraphs")(fn)
-        self.assertEqual(counters["inductor"]["select_algorithm_precompile"], 0)
-
     def test_cat_addmm(self):
         def fn(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor):
             return torch.cat(
diff --git a/test/inductor/test_select_algorithm.py b/test/inductor/test_select_algorithm.py
@@ -19,10 +19,8 @@
 
 
 def patches(fn):
-    def skip_cache(self, choices, name, key, benchmark):
-        if benchmark is None:
-            return {}
-        return benchmark(choices)
+    def skip_cache(self, choices, name, key, generate):
+        return generate(choices)
 
     for patcher in [
         dynamo_config.patch(verbose=True),
diff --git a/torch/_inductor/autotune_process.py b/torch/_inductor/autotune_process.py
@@ -502,6 +502,7 @@ def benchmark(
 class TritonBenchmarkRequest(BenchmarkRequest):
     # Important: Instances of this class have to be serializable
     # across process boundaries. Do not put CUDA Tensors in here!
+
     def __init__(
         self,
         kernel_name: str,
@@ -544,8 +545,6 @@ def make_run_fn(
         if "warmup" in inspect.signature(run_method).parameters:
             warmup_arg["warmup"] = False
 
-        from torch._C import _cuda_getCurrentRawStream as get_raw_stream
-
         if torch.version.hip and self.matrix_instr_nonkdim != 0:
             return functools.partial(
                 run_method,
@@ -554,7 +553,9 @@ def make_run_fn(
                 *self.extra_args,
                 grid=self.grid,
                 **warmup_arg,
-                stream=get_raw_stream(self.output_tensor_meta.device.index),
+                num_stages=self.num_stages,
+                num_warps=self.num_warps,
+                matrix_instr_nonkdim=self.matrix_instr_nonkdim,
             )
         else:
             return functools.partial(
@@ -564,13 +565,10 @@ def make_run_fn(
                 *self.extra_args,
                 grid=self.grid,
                 **warmup_arg,
-                stream=get_raw_stream(self.output_tensor_meta.device.index),
+                num_stages=self.num_stages,
+                num_warps=self.num_warps,
             )
 
-    def precompile(self):
-        mod = PyCodeCache.load_by_key_path(self.module_cache_key, self.module_path)
-        getattr(mod, self.kernel_name).precompile()
-
     def __str__(self) -> str:
         return f"{self.kernel_name=}, {self.module_path=}, {self.module_cache_key=}"
 
diff --git a/torch/_inductor/codecache.py b/torch/_inductor/codecache.py
@@ -156,7 +156,7 @@ def get_system() -> Dict[str, Any]:
         try:
             import triton
 
-            triton_version = triton.__version__  # type: ignore[attr-defined]
+            triton_version = triton.__version__
         except ModuleNotFoundError:
             triton_version = None
 
@@ -262,15 +262,15 @@ def lookup(
         choices: List[ChoiceCaller],
         op: str,
         inputs: str,
-        benchmark: Optional[Callable[[Any], Dict[ChoiceCaller, float]]],
+        benchmark: Callable[[Any], Dict[ChoiceCaller, float]],
     ) -> Dict[ChoiceCaller, float]:
         """
         Check to see if we have benchmarked the given choice callers. For each
         choice caller:
 
             1. Check global_cache[op][inputs][choice][precision], return benchmark if cached.
             2. Check local_cache[op][inputs][choice][precision], return benchmark if cached.
-            3. If benchmark is not None:
+            3.
                 a. `max_autotune_gemm=True`: benchmark the choice, update
                     local_cache[op][inputs][choice], and return the benchmark.
                 b. `max_autotune_gemm=False`: don't benchmark the choice, return nothing.
@@ -303,13 +303,9 @@ def check_cache(cache, callback=None) -> bool:
         if config.max_autotune or config.max_autotune_gemm:
             local_cache = self.get_local_cache()
             # check local cache first since it is data specific to the current machine
-            if (
-                not check_cache(local_cache)
-                and not (
-                    use_global_cache()
-                    and check_cache(self.get_global_cache(), callback=log_stats)
-                )
-                and benchmark is not None
+            if not check_cache(local_cache) and not (
+                use_global_cache()
+                and check_cache(self.get_global_cache(), callback=log_stats)
             ):
                 try:
                     # re-benchmark everything to try to get consistent numbers from the same machine
diff --git a/torch/_inductor/codegen/triton_utils.py b/torch/_inductor/codegen/triton_utils.py
@@ -65,32 +65,6 @@ def signature_to_meta(
     }
 
 
-def is_unaligned_buffer(arg: TensorArg):
-    buf_name = arg.buffer
-    if buf_name in V.graph.graph_inputs:
-        return not config.assume_aligned_inputs
-
-    if buf_name in V.graph.constants:
-        # all constants are assumed to be aligned
-        return False
-
-    if V.graph.scheduler:
-        layout = V.graph.scheduler.get_buffer_layout(buf_name)
-    else:
-        buffer = V.graph.get_buffer(buf_name)
-        # output arg
-        if not buffer:
-            assert buf_name == V.kernel.output_node.name
-            layout = V.kernel.output_node.layout
-        else:
-            layout = buffer.get_layout()
-
-    if isinstance(layout, torch._inductor.ir.NonOwningLayout):
-        return not layout.maybe_guard_aligned()
-    else:
-        return False
-
-
 def config_of(
     args: List[KernelArgType],
     *,
@@ -109,7 +83,9 @@ def is_aligned(x: KernelArgType, alignment: int, include_tensor: bool) -> bool:
                 offset_aligned = V.graph.sizevars.statically_known_multiple_of(
                     x.offset * x.dtype.itemsize, alignment  # type: ignore[arg-type]
                 )
-                return offset_aligned and not is_unaligned_buffer(x)
+                return offset_aligned and not V.graph.scheduler.is_unaligned_buffer(
+                    x.buffer
+                )
             else:
                 return False
         if isinstance(x, SizeArg):
diff --git a/torch/_inductor/scheduler.py b/torch/_inductor/scheduler.py
@@ -2495,9 +2495,18 @@ def codegen(self):
 
         self.flush()
 
-    def get_buffer_layout(self, buf_name: str) -> ir.Layout:
+    def is_unaligned_buffer(self, buf_name):
+        if buf_name in V.graph.graph_inputs:
+            return not config.assume_aligned_inputs
+        if buf_name in V.graph.constants:
+            # all constants are assumed to be aligned
+            return False
         node = self.name_to_node[buf_name]
-        return node.node.get_layout()
+        layout = node.node.get_layout()
+        if isinstance(layout, ir.NonOwningLayout):
+            return not layout.maybe_guard_aligned()
+        else:
+            return False
 
 
 class BaseScheduling:
diff --git a/torch/_inductor/select_algorithm.py b/torch/_inductor/select_algorithm.py
@@ -94,7 +94,7 @@ def __init__(
         grid_fn,
         meta,
         call_sizes,
-        use_jit=False,
+        use_jit=True,
         prefix_args=0,
         suffix_args=0,
         epilogue_fn=identity,
@@ -153,8 +153,8 @@ def jit_lines(self):
         argdefs, _, signature = self.args.python_argdefs()
         triton_meta = {
             "signature": signature_to_meta(signature, size_dtype=self.index_dtype),
-            "device": self.output_node.get_device().index,
-            "device_type": self.output_node.get_device().type,
+            "device": V.graph.scheduler.current_device.index,
+            "device_type": V.graph.scheduler.current_device.type,
             "constants": {},
         }
         triton_meta["configs"] = [config_of(signature)]
@@ -554,7 +554,7 @@ def generate(
         ), TritonTemplateKernel(
             kernel_name=kernel_name,
             output_node=fake_out,
-            use_jit=False,
+            use_jit=True,
             **kernel_options,
         ) as kernel:
             try:
@@ -740,10 +740,6 @@ def benchmark(self, *args, out):
         assert self.bmreq is not None
         return self.bmreq.benchmark(*args, output_tensor=out)
 
-    def precompile(self):
-        assert self.bmreq is not None
-        self.bmreq.precompile()
-
     def __str__(self):
         return f"TritonTemplateCaller({self.bmreq.module_path}, {self.debug_extra})"
 
@@ -885,7 +881,6 @@ def __call__(
 
         # TODO(nmacchioni): remove once CI tests are fixed
         choices = [choice for choice in choices if choice is not None]
-
         if len(choices) == 0:
             raise RuntimeError(
                 "No choices to select, please consider adding ATEN into max_autotune_gemm_backends "
@@ -902,38 +897,19 @@ def __call__(
         def make_benchmark_fn():
             return self.make_benchmark_fn(choices, input_nodes, layout, input_gen_fns)
 
-        def precompile(choices) -> Callable[[], None]:
-            def no_op(*args, **kwargs):
-                return
-
+        def precompile(choices):
             if (
                 precompilation_timeout_seconds is None
                 or precompilation_timeout_seconds <= 0
             ):
-                return no_op
+                return
             num_workers = min(
                 config.compile_threads,
                 torch.get_num_threads(),
                 len(choices),
             )
             if num_workers <= 0:
-                return no_op
-
-            # TODO - debug issue
-            if torch.version.hip:
-                return no_op
-
-            # check local and global cache before precompiling
-            timings = self.lookup(
-                choices,
-                name,
-                repr([self.key_of(x) for x in input_nodes]),
-                benchmark=None,
-            )
-
-            if timings:
-                return no_op
-
+                return
             log.info(
                 "Multithreaded precompilation for %d choices using %d worker threads",
                 len(choices),
