add *BLOCK like constants to launcher args

Signed-off-by: Anatoly Myachev <[email protected]>

Author: anmyachev

scripts/patch-pytorch.sh

@@ -17,5 +17,5 @@ echo "Applying PyTorch patches in $REPO_ROOT"
 cd "$REPO_ROOT"
 
 curl -sSL https://github.com/pytorch/pytorch/pull/126516.diff | git apply -
-# trigger 3
+# trigger 4
 git apply "${SCRIPT_DIR}/pytorch.patch"

scripts/pytorch.patch

@@ -110,6 +110,24 @@ index 00031a56b8d..59086d41b40 100644
          self.triton_meta = triton_meta
 
          for tree in self.range_trees:
+diff --git a/torch/_inductor/codegen/triton_utils.py b/torch/_inductor/codegen/triton_utils.py
+index 8b8c29bbb15..c89a76e9868 100644
+--- a/torch/_inductor/codegen/triton_utils.py
++++ b/torch/_inductor/codegen/triton_utils.py
+@@ -165,12 +165,4 @@ def config_of(
+     else:
+         divisible_by_16 = ()
+ 
+-    equal_to_1 = tuple(
+-        i
+-        for i, arg in zip(indices, args)
+-        if isinstance(arg, SizeArg)
+-        and isinstance(arg.expr, (int, sympy.Integer))
+-        and V.graph.sizevars.statically_known_equals(arg.expr, 1)  # type: ignore[arg-type]
+-    )
+-
+-    return AttrsDescriptorWrapper(divisible_by_16, equal_to_1)
++    return AttrsDescriptorWrapper(divisible_by_16)
 diff --git a/torch/_inductor/codegen/wrapper.py b/torch/_inductor/codegen/wrapper.py
 index 2ab2b326354..42d76b8bf94 100644
 --- a/torch/_inductor/codegen/wrapper.py
@@ -123,25 +141,24 @@ index 2ab2b326354..42d76b8bf94 100644
              "configs": [
                  config_of(
 diff --git a/torch/_inductor/runtime/hints.py b/torch/_inductor/runtime/hints.py
-index fa2a1334380..11de4dc6e5f 100644
+index fa2a1334380..4d730fd45de 100644
 --- a/torch/_inductor/runtime/hints.py
 +++ b/torch/_inductor/runtime/hints.py
-@@ -44,25 +44,16 @@ def _is_triton_available() -> bool:
+@@ -44,25 +44,14 @@ def _is_triton_available() -> bool:
  # Define `AttrsDescriptorWrapper` function with clear conditional handling
  if _is_triton_available():
      try:
 -        from triton.backends.compiler import AttrsDescriptor
 
          def AttrsDescriptorWrapper(
              divisible_by_16=None,
-             equal_to_1=None,
+-            equal_to_1=None,
          ):
 -            # Prepare the arguments for AttrsDescriptor
              kwargs = {
 -                "tt.divisibility": divisible_by_16,
 -                "tt.equal_to": equal_to_1,
 +                tuple([(i,) for i in divisible_by_16]): [["tt.divisibility", 16]],
-+                tuple([(i,) for i in equal_to_1]): [["tt.equal_to", 1]],
              }
 -
 -            # Instantiate AttrsDescriptor with the prepared arguments
@@ -156,7 +173,7 @@ index fa2a1334380..11de4dc6e5f 100644
      except ImportError:
          from triton.compiler.compiler import AttrsDescriptor
 diff --git a/torch/_inductor/runtime/triton_heuristics.py b/torch/_inductor/runtime/triton_heuristics.py
-index 281d0e78ba4..8e7ca6b5822 100644
+index 281d0e78ba4..3b059a365c9 100644
 --- a/torch/_inductor/runtime/triton_heuristics.py
 +++ b/torch/_inductor/runtime/triton_heuristics.py
 @@ -414,10 +414,21 @@ class CachingAutotuner(KernelInterface):
@@ -196,3 +213,19 @@ index 281d0e78ba4..8e7ca6b5822 100644
          ]
          binary_shared = (
              binary.shared if hasattr(binary, "shared") else binary.metadata.shared
+@@ -952,6 +961,7 @@ class CachingAutotuner(KernelInterface):
+             ):
+                 return launcher(
+                     *args,
++                    **launcher.config.kwargs,
+                     **kwargs,
+                     grid=grid,
+                     stream=stream,
+@@ -959,6 +969,7 @@ class CachingAutotuner(KernelInterface):
+         else:
+             return launcher(
+                 *args,
++                **launcher.config.kwargs,
+                 **kwargs,
+                 grid=grid,
+                 stream=stream,