Add some Triton related suppressions that don't show on CI (pytorch#166868)

Signed-off-by: Edward Z. Yang <[email protected]>

Pull Request resolved: pytorch#166868
Approved by: https://github.com/maggiemoss, https://github.com/zou3519

Author: ezyang

torch/_dynamo/repro/after_aot.py

@@ -405,6 +405,7 @@ def generate_compiler_repro_string(
                 # pyrefly: ignore [missing-attribute]
                 kernel._fn_name
                 if isinstance(kernel, JITFunction)
+                # pyrefly: ignore  # missing-attribute
                 else kernel.fn._fn_name
             )
             fn_name = fn_name.split(".")[-1]

torch/_higher_order_ops/triton_kernel_wrap.py

@@ -264,6 +264,7 @@ def generate_ttir(
 
     assert isinstance(kernel, JITFunction)
 
+    # pyrefly: ignore  # missing-attribute
     context = triton._C.libtriton.ir.context()
     target = triton.runtime.driver.active.get_current_target()
     backend = triton.compiler.compiler.make_backend(target)
@@ -305,6 +306,7 @@ def generate_ttir(
                 base_tensor = torch.empty(
                     [elements_per_dim] * len(block_shape), dtype=a.dtype
                 )
+            # pyrefly: ignore  # bad-argument-type
             ordered_args[name] = TensorDescriptor.from_tensor(base_tensor, block_shape)
         elif isinstance(a, (FakeTensor, torch._inductor.ir.TensorBox)):
             with torch._C._DisableTorchDispatch():
@@ -368,6 +370,7 @@ def _get_specialization(args):  # type: ignore[no-untyped-def]
 
             target = triton.runtime.driver.active.get_current_target()
             backend_ = triton.compiler.compiler.make_backend(target)
+            # pyrefly: ignore  # missing-attribute
             return backend_.get_attrs_descriptor(args, kernel.params)
         else:
             assert (
@@ -384,6 +387,7 @@ def _get_specialization(args):  # type: ignore[no-untyped-def]
                 except TypeError:  # Unknown arg `specialize_extra`
                     # Older versions of Triton take specialize_extra as an arg to specialize_impl
                     specialize_impl = functools.partial(
+                        # pyrefly: ignore  # missing-argument
                         triton.runtime.jit.create_specialize_impl(),
                         specialize_extra=backend.get_arg_specialization,
                     )
@@ -468,6 +472,7 @@ def get_signature_value(idx: int, arg: Any) -> str:
             if i not in constexprs
         }
 
+    # pyrefly: ignore  # missing-attribute
     triton._C.libtriton.ir.load_dialects(context)
     backend.load_dialects(context)
 
@@ -477,22 +482,29 @@ def get_signature_value(idx: int, arg: Any) -> str:
     # backward compatibility here.
     make_ir_sig_params = len(inspect.signature(src.make_ir).parameters)
     get_codegen_implementation_sig_params = len(
+        # pyrefly: ignore  # missing-attribute
         inspect.signature(backend.get_codegen_implementation).parameters
     )
     if make_ir_sig_params == 2:
+        # pyrefly: ignore  # missing-argument
         ttir_module = src.make_ir(options, context)
     elif make_ir_sig_params == 3:
+        # pyrefly: ignore  # missing-attribute
         codegen_fns = backend.get_codegen_implementation()
+        # pyrefly: ignore  # missing-argument
         ttir_module = src.make_ir(options, codegen_fns, context)
     elif make_ir_sig_params == 4:
         codegen_args = [options] if get_codegen_implementation_sig_params == 1 else []
+        # pyrefly: ignore  # missing-attribute
         codegen_fns = backend.get_codegen_implementation(*codegen_args)
         module_map = backend.get_module_map()
         ttir_module = src.make_ir(options, codegen_fns, module_map, context)
     else:
         codegen_args = [options] if get_codegen_implementation_sig_params == 1 else []
+        # pyrefly: ignore  # missing-attribute
         codegen_fns = backend.get_codegen_implementation(*codegen_args)
         module_map = backend.get_module_map()
+        # pyrefly: ignore  # bad-argument-count
         ttir_module = src.make_ir(target, options, codegen_fns, module_map, context)
     if not ttir_module.verify():
         raise RuntimeError("Verification for TTIR module has failed")
@@ -1102,6 +1114,7 @@ def triton_kernel_wrapper_mutation_dense(
                 from triton.tools.tensor_descriptor import TensorDescriptor
 
                 block_shape = stable_meta[0]
+                # pyrefly: ignore  # bad-argument-type
                 kwargs[k] = TensorDescriptor.from_tensor(tensor, block_shape)
 
     # move as many positional arguments from dicts to args as we
@@ -1658,6 +1671,7 @@ def call_triton_kernel(
                     "Passing multiple @triton.autotune decorators is not supported. "
                     "Please use a single @triton.autotune decorator instead."
                 )
+            # pyrefly: ignore  # missing-attribute
             iter_kernel = iter_kernel.fn
 
         # Process the @triton.heuristics decorator:
@@ -1868,6 +1882,7 @@ def call_triton_kernel(
 
         # Both for grid's meta as well as for the kernel, we need combined
         # args and kwargs combined and normalized
+        # pyrefly: ignore  # missing-attribute
         combined_args_raw = {**dict(zip(variable.kernel.arg_names, args)), **kwargs}
 
         # precompute the grid for the kernel
@@ -2061,6 +2076,7 @@ def __init__(
         kernel_idx: Optional[int],
         grid: Optional["TritonGridType"],
     ) -> None:
+        # pyrefly: ignore  # bad-assignment
         self.kernel = None
         self.grid = None
         tracing_triton_hopifier_singleton.init_variable(self, kernel, kernel_idx, grid)

torch/_inductor/codegen/wrapper.py

@@ -289,11 +289,15 @@ def traverse(cur_kernel):
                 if isinstance(symbol, JITFunction):
                     compile_wrapper.newline()
                     compile_wrapper.writeline("@triton.jit")
+                    # pyrefly: ignore  # missing-attribute
                     compile_wrapper.splice(symbol.src, strip=True)
                     symbols_included.add(symbol_name)
                     traverse(symbol)
                 elif hasattr(triton, "constexpr_function") and isinstance(
-                    symbol, triton.runtime.jit.ConstexprFunction
+                    # pyrefly: ignore  # missing-attribute
+                    symbol,
+                    # pyrefly: ignore  # missing-attribute
+                    triton.runtime.jit.ConstexprFunction,
                 ):
                     compile_wrapper.newline()
                     compile_wrapper.writeline("@triton.constexpr_function")

torch/_inductor/codegen/wrapper_fxir.py

@@ -949,7 +949,9 @@ def tune_kernel(tuner: CachingAutotuner, call_args: Sequence[Any]) -> None:
             from triton.runtime import driver
 
             log.info("Autotuning Triton kernel %s at compile time.", kernel_name)
+            # pyrefly: ignore  # missing-attribute
             device = driver.active.get_current_device()
+            # pyrefly: ignore  # missing-attribute
             stream = driver.active.get_current_stream(device)
 
             def node_to_tuning_arg(arg: Any) -> Any:

torch/_inductor/ir.py

@@ -6985,6 +6985,7 @@ def get_kernel_and_metadata(self) -> tuple[Kernel, Any, list[str], list[str]]:
 
             configs = kernel.configs
             kernel = kernel.fn
+        # pyrefly: ignore  # bad-return
         return kernel, configs, restore_value_args, reset_to_zero_args
 
     @override
@@ -7140,7 +7141,10 @@ def __init__(
         self.mutable_args = [
             kernel_args[key]
             for key in identify_mutated_tensors(
-                kernel, {**kernel_args, **autotuned_kwargs}, tma_descriptor_metadata
+                # pyrefly: ignore  # bad-argument-type
+                kernel,
+                {**kernel_args, **autotuned_kwargs},
+                tma_descriptor_metadata,
             )
         ]
 

torch/_inductor/utils.py

@@ -2555,11 +2555,14 @@ def get_device_tflops(dtype: torch.dtype) -> float:
             return get_max_simd_tflops(torch.float32, sm_clock)
     else:
         if dtype in (torch.float16, torch.bfloat16) and SM80OrLater:
+            # pyrefly: ignore  # missing-argument
             return get_max_tensorcore_tflops(dtype)
 
         if torch.backends.cuda.matmul.allow_tf32:
+            # pyrefly: ignore  # missing-argument
             return get_max_tensorcore_tflops(torch.float32)
         else:
+            # pyrefly: ignore  # missing-argument
             return get_max_simd_tflops(torch.float32)
 
 
@@ -2573,6 +2576,7 @@ def get_gpu_dram_gbps() -> int:
 def get_gpu_shared_memory() -> int:
     from triton.runtime import driver
 
+    # pyrefly: ignore  # missing-attribute
     return driver.active.utils.get_device_properties(0).get("max_shared_mem", 0)
 
 

torch/sparse/_triton_ops.py

@@ -1302,17 +1302,28 @@ def kernel(grid, *sliced_tensors):
         # pyrefly: ignore [unsupported-operation]
         _bsr_strided_addmm_kernel[grid](
             *ptr_stride_extractor(*sliced_tensors),
+            # pyrefly: ignore  # bad-argument-count
             beta,
             alpha,
+            # pyrefly: ignore  # bad-keyword-argument, bad-argument-type
             beta_is_one=beta == 1,
+            # pyrefly: ignore  # bad-keyword-argument, bad-argument-type
             beta_is_nonzero=beta != 0,
+            # pyrefly: ignore  # bad-keyword-argument, bad-argument-type
             alpha_is_one=alpha == 1,
+            # pyrefly: ignore  # bad-keyword-argument, bad-argument-type
             left_alpha_is_one=left_alpha_is_one,
+            # pyrefly: ignore  # bad-keyword-argument, bad-argument-type
             right_alpha_is_one=right_alpha_is_one,
+            # pyrefly: ignore  # bad-keyword-argument, bad-argument-type
             BLOCKSIZE_ROW=BM,
+            # pyrefly: ignore  # bad-keyword-argument, bad-argument-type
             BLOCKSIZE_INNER=BK,
+            # pyrefly: ignore  # bad-keyword-argument
             BLOCKSIZE_COL=BN,
+            # pyrefly: ignore  # bad-keyword-argument
             allow_tf32=dot_out_dtype == tl.float32,
+            # pyrefly: ignore  # bad-keyword-argument, bad-argument-type
             acc_dtype=dot_out_dtype,
             **meta,
         )
@@ -1633,12 +1644,17 @@ def kernel(grid, *sliced_tensors):
                 beta,
                 is_beta_zero,
                 *blocksize,
+                # pyrefly: ignore  # bad-argument-count
                 k,
                 tile_k,
                 *ptr_stride_extractor(*sliced_tensors),
+                # pyrefly: ignore  # bad-keyword-argument, bad-argument-type
                 acc_dtype=acc_dtype,
+                # pyrefly: ignore  # bad-keyword-argument, bad-argument-type
                 allow_tf32=allow_tf32,
+                # pyrefly: ignore  # unexpected-keyword
                 num_stages=1,
+                # pyrefly: ignore  # unexpected-keyword
                 num_warps=4,
             )
 
@@ -1923,6 +1939,7 @@ def bsr_softmax(input, max_row_nnz=None):
         def kernel(grid, *sliced_tensors):
             _bsr_softmax_kernel[grid](
                 *ptr_stride_extractor(*sliced_tensors),
+                # pyrefly: ignore  # bad-argument-count
                 row_block,
                 col_block,
                 max_row_nnz,
@@ -2096,8 +2113,11 @@ def grid(META):
         if "allow_tf32" not in meta:
             meta.update(allow_tf32=dot_out_dtype == tl.float32)
         _scatter_mm2_kernel[grid](
+            # pyrefly: ignore  # bad-argument-type
             M,
+            # pyrefly: ignore  # bad-argument-type
             K,
+            # pyrefly: ignore  # bad-argument-type
             N,
             blocks,
             blocks.stride(0),
@@ -2116,7 +2136,9 @@ def grid(META):
             pq_indices,
             pq_indices.stride(0),
             pq_indices.stride(1),
+            # pyrefly: ignore  # bad-argument-type
             dot_out_dtype=dot_out_dtype,
+            # pyrefly: ignore  # bad-argument-type
             **meta,
         )
 
@@ -2299,6 +2321,7 @@ def grid(META):
         _scatter_mm6_kernel[grid](
             B,
             Ms,
+            # pyrefly: ignore  # bad-argument-type
             Ks,
             N,
             blocks,
@@ -2317,6 +2340,7 @@ def grid(META):
             r_offsets,
             p_offsets,
             q_offsets,
+            # pyrefly: ignore  # bad-argument-type
             dot_out_dtype=dot_out_dtype,
             **meta,
         )