[inductor] More fixes on the keys of constants and signature dictionaries (pytorch#135406)

Previous PR forgets to change two other places that also create `constants` and `signature`. pytorch#135170 Pull Request resolved: pytorch#135406 Approved by: https://github.com/jansel

Author: Jokeren

docs/source/torch.compiler_get_started.rst

@@ -57,7 +57,7 @@ the following:
 
 .. code-block:: python
 
- @pointwise(size_hints=[16384], filename=__file__, triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: 'i32'}, 'device': 0, 'constants': {}, 'mutated_arg_names': [], 'configs': [instance_descriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]})
+ @pointwise(size_hints=[16384], filename=__file__, triton_meta={'signature': {'in_ptr0': '*fp32', 'out_ptr0': '*fp32', 'xnumel': 'i32'}, 'device': 0, 'constants': {}, 'mutated_arg_names': [], 'configs': [instance_descriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]})
  @triton.jit
  def triton_(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
  xnumel = 10000

test/inductor/test_cuda_repro.py

@@ -433,7 +433,11 @@ def decorator(fn):
  triton.Config({"XBLOCK": 2}),
  ],
  meta={
- "signature": {0: "*fp32", 1: "*fp32", 2: "i32"},
+ "signature": {
+ "in_out_ptr0": "*fp32",
+ "in_ptr0": "*fp32",
+ "xnumel": "i32",
+ },
  "device": DeviceProperties.create(torch.device("cuda")),
  "configs": [instance_descriptor(divisible_by_16=(0, 1), equal_to_1=())],
  "constants": {},

test/inductor/test_metrics.py

@@ -14,7 +14,7 @@
  reduction_hint=ReductionHint.INNER,
  filename=__file__,
  triton_meta={
- 'signature': {0: '*fp32', 1: '*fp32', 2: 'i32', 3: 'i32'},
+ 'signature': {'in_out_ptr0': '*fp32', 'in_ptr0': '*fp32', 'xnumel': 'i32', 'rnumel': 'i32'},
  'device': 0,
  'device_type': 'GPU_TYPE',
  'constants': {},

test/inductor/test_triton_heuristics.py

@@ -102,7 +102,7 @@ def triton_(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
  tl.store(out_ptr0 + (x0), tmp1, xmask)
 
  triton_meta = {
- "signature": {0: "*fp32", 1: "*fp32", 2: "i32"},
+ "signature": {"in_ptr0": "*fp32", "out_ptr0": "*fp32", "xnumel": "i32"},
  "device": DeviceProperties.create(torch.device("cuda")),
  "constants": {},
  "configs": [AttrsDescriptor(divisible_by_16=(0, 1, 2), equal_to_1=())],

torch/_inductor/codegen/triton.py

@@ -2648,7 +2648,7 @@ def codegen_kernel(self, name=None):
  mutated_args = sorted(mutated_args)
 
  triton_meta_signature = signature_to_meta(
- signature, size_dtype=self.index_dtype
+ signature, size_dtype=self.index_dtype, argdefs=argdefs
  )
  triton_meta = {
  "signature": triton_meta_signature,
@@ -2676,7 +2676,7 @@ def codegen_kernel(self, name=None):
  for tree in self.active_range_trees():
  sizearg = SizeArg(f"{tree.prefix}numel", tree.numel)
  signature.append(sizearg)
- triton_meta_signature[len(argdefs)] = signature_of(
+ triton_meta_signature[sizearg.name] = signature_of(
  sizearg, size_dtype=self.index_dtype
  )
  argdefs.append(f"{tree.prefix}numel")
@@ -2694,7 +2694,7 @@ def codegen_kernel(self, name=None):
  # https://github.com/pytorch/pytorch/issues/120478#issuecomment-1962822307
  # https://github.com/openai/triton/blob/231efe9ed2d200be0f69a07c298e4342b08efe3d/python/triton/runtime/jit.py#L384
  for arg_num in triton_meta["configs"][0].equal_to_1: # type: ignore[index]
- triton_meta["constants"][arg_num] = 1 # type: ignore[index]
+ triton_meta["constants"][signature[arg_num].name] = 1 # type: ignore[index]
 
  self.triton_meta = triton_meta
 

torch/_inductor/codegen/triton_combo_kernel.py

@@ -660,18 +660,19 @@ def jit_line(
  heuristics: str,
  size_hints: List[int],
  selected_kernel: TritonKernel,
+ signature: List[Any],
+ argdefs: List[str],
  pointwise_with_reduce: bool = False,
- signature: Optional[List[Any]] = None,
  ) -> str:
  can_use_32bit = all(k.index_dtype == "tl.int32" for k in self.sub_kernels)
  size_dtype = "tl.int32" if can_use_32bit else "tl.int64"
- if signature is None:
- _, _, signature, _ = self.args.python_argdefs()
  for i, sub in enumerate(self.sub_kernels):
  self.min_x_blocks_sub_kernel(sub, i)
  self.select_dispatch_strategy()
  triton_meta = {
- "signature": signature_to_meta(signature, size_dtype=size_dtype),
+ "signature": signature_to_meta(
+ signature, size_dtype=size_dtype, argdefs=argdefs
+ ),
  "device": DeviceProperties.create(
  V.graph.scheduler.get_current_device_or_throw()
  ),
@@ -850,6 +851,7 @@ def codegen_kernel(self, name: Optional[str] = None) -> str:
  selected_kernel,
  pointwise_with_reduce=pointwise_with_reduction,
  signature=signature,
+ argdefs=argdefs,
  )
  )
  code.writeline(

torch/_inductor/codegen/triton_utils.py

@@ -68,12 +68,13 @@ def signature_to_meta(
  signature: List[KernelArgType],
  *,
  size_dtype: str,
+ argdefs: List[str],
  indices: Optional[List[int]] = None,
-) -> Dict[int, str]:
+) -> Dict[str, str]:
  if indices is None:
  indices = list(range(len(signature)))
  return {
- i: signature_of(arg, size_dtype=size_dtype)
+ argdefs[i]: signature_of(arg, size_dtype=size_dtype)
  for i, arg in zip(indices, signature)
  }
 

torch/_inductor/codegen/wrapper.py

@@ -1275,15 +1275,15 @@ def define_user_defined_triton_kernel(self, kernel, configs, kwargs):
  from .common import KernelArgType, SizeArg, TensorArg
 
  signature: List[KernelArgType] = []
- constants: Dict[int, Any] = {}
+ constants: Dict[str, Any] = {}
  non_constant_indices = []
- equal_to_1_arg_idx: List[int] = []
+ equal_to_1_args: List[str] = []
  for idx, key in enumerate(kernel.arg_names):
  if key not in kwargs:
  continue
  arg = kwargs[key]
  if idx in kernel.constexprs:
- constants[idx] = arg
+ constants[key] = arg
  else:
  non_constant_indices.append(idx)
  if isinstance(arg, ir.Buffer):
@@ -1313,13 +1313,14 @@ def define_user_defined_triton_kernel(self, kernel, configs, kwargs):
  ) and V.graph.sizevars.statically_known_equals(
  arg, 1 # type: ignore[arg-type]
  ):
- equal_to_1_arg_idx.append(idx)
+ equal_to_1_args.append(key)
  index_dtype = "tl.int32"
  triton_meta = {
  "signature": signature_to_meta(
  signature,
  size_dtype=index_dtype,
  indices=non_constant_indices,
+ argdefs=kernel.arg_names,
  ),
  "device": DeviceProperties.create(
  V.graph.scheduler.get_current_device_or_throw()
@@ -1333,7 +1334,7 @@ def define_user_defined_triton_kernel(self, kernel, configs, kwargs):
  # https://github.com/openai/triton/blob/231efe9ed2d200be0f69a07c298e4342b08efe3d/python/triton/runtime/jit.py#L384
  "constants": {
  **constants,
- **dict.fromkeys(equal_to_1_arg_idx, 1),
+ **dict.fromkeys(equal_to_1_args, 1),
  },
  "configs": [
  config_of(

torch/_inductor/runtime/triton_heuristics.py

@@ -359,7 +359,7 @@ def _precompile_config(self, cfg: Config, warm_cache_only: bool):
  if k == "waves_per_eu":
  compile_meta["waves_per_eu"] = v
  continue
- compile_meta["constants"][self.fn.arg_names.index(k)] = v
+ compile_meta["constants"][k] = v
  compile_meta["num_warps"] = cfg.num_warps
  compile_meta["num_stages"] = cfg.num_stages
  compile_meta["debug"] = self.inductor_meta.get(

torch/_inductor/select_algorithm.py

@@ -214,13 +214,15 @@ def jit_lines(self):
 
  argdefs, _, signature, _ = self.args.python_argdefs()
  triton_meta = {
- "signature": signature_to_meta(signature, size_dtype=self.index_dtype),
+ "signature": signature_to_meta(
+ signature, size_dtype=self.index_dtype, argdefs=argdefs
+ ),
  "device": DeviceProperties.create(self.output_node.get_device()),
  "constants": {},
  }
  triton_meta["configs"] = [config_of(signature)]
  for arg_num in triton_meta["configs"][0].equal_to_1: # type: ignore[index]
- triton_meta["constants"][arg_num] = 1 # type: ignore[index]
+ triton_meta["constants"][signature[arg_num].name] = 1 # type: ignore[index]
  matrix_instr_nonkdim = self.meta.get("matrix_instr_nonkdim", 0)
  if matrix_instr_nonkdim != 0:
  triton_meta["matrix_instr_nonkdim"] = matrix_instr_nonkdim