[inductor] Replace `Any` with precise types in symbolic-shape predicates (#194482)

Inductor threads a heterogeneous "maybe-symbolic scalar" through its shape
reasoning -- a dimension that may be a plain `int`, a `sympy.Expr`, a `SymInt`,
or an `fx.Node` carrying one of those in `meta["val"]`. The small predicate
helpers that consume it were all annotated `Any`, so none of their call sites
were checked at all.

Reading them, they are not all the same shape, and this PR gives each the type
that actually fits rather than applying one blanket replacement:

- Helpers whose body `isinstance`-branches over every case take `object`. That
  matches the four helpers of exactly this shape that already use `object`:
  `ir._is_static`, `codegen/cutlass/utils.is_static_int`,
  `codegen/triton.is_sympy_integer_like`, and
  `codegen/cutedsl/cutedsl_kernel._static_integer_expr`.
- Helpers whose body does arithmetic or ordering on the value take `_IntLike`
  (`int | sympy.Expr`). `object` would be wrong there: `%`, `*`, and `<=` are
  not defined on `object`, and reaching for a `cast` to keep the `object`
  annotation would have thrown away the checking the annotation is there to buy.
- Helpers whose body indexes and iterates a size take `Sequence[_IntLike]`, and
  the scaled-mm heuristic's scale arguments -- on which the body calls
  `get_size()` -- take `IRNode`. The one exception is
  `statically_known_shape_equal`, which is handed `FakeTensor.shape` and
  `.stride()`, so its elements can be `SymInt` rather than `sympy.Expr`; since
  it only forwards them to `statically_known_equal`, its element type is
  `object`.

Narrowing a predicate buys nothing if its own caller still launders the value
through `Any` first, so the callers are narrowed up to the frame where the value
enters: `is_flex_gemm_partial_reduction_shape` and
`flex_gemm_output_config_supported` in `kernel/flex_gemm/constraints.py`,
`validate_flex_gemm_aux_outputs` and `flex_gemm_config_keys` in
`kernel/flex_gemm/lowering.py`, and the `NodeInfo` fields that
`codegen/triton_combo_kernel._size_hint` reads (`codegen/simd.py`).

Doing that surfaced several annotations that were already wrong, not merely
loose:

- `statically_known_multiple`'s `divisor` was `int`, but both worlds reach it:
  a plain `int` from `group` and `swap_ab_alignment`, and a `sympy.Expr` from
  `aux_m` / `aux_n`, which originate in `ir.convert_shape_to_inductor`
  (`-> list[sympy.Expr]`). `_IntLike` is the only type covering all four call
  sites.
- `flex_gemm_config_keys` declared `m: int, n: int`. Its production caller
  passes `TensorBox.get_size()` elements, and `test_flex_gemm.py:1169` already
  passes a `sympy.Symbol` for `n`. The function's own `has_free_symbols((n,))`
  guard and the `guard_or_false(sympy.Eq(n % swap_ab_alignment, 0))` branch
  below it would both be dead code if `n` really were an `int`.

`kernel/flex_gemm/constraints.py` needs two unions rather than one, neither of
which is `_IntLike` alone:

- Its local-reduce validators
  (`validate_local_reduce_selected_dim_divisible`,
  `local_reduce_compressed_shape`) are reached only with `torch.Size`-derived
  tuples -- from real tensors at `runtime.py:306, 363, 520`, and from
  `tuple(FakeTensor.shape)` at `fx_cutedsl_codegen.py:271` ->
  `quack_reductions.py:246`. So they take `Sequence[IntLikeType]`
  (`int | SymInt`). Keeping `sympy.Expr` out of them is not cosmetic: sympy's
  `==`/`!=` are structural, so `Mod(s0, 8) != 0` is Python `True`, and a
  merely-unknown dimension would raise `LOCAL_REDUCE_DIVISIBLE_SHAPE_ERROR`
  rather than fall through the way `constraints.py:239-242` intends.
- `statically_known_multiple` sits below both those validators and the
  inductor-lowering callers, so it takes `_IntLike | IntLikeType`. `object` is
  not available for either, since the bodies need `%` and `//`.

`NodeInfo.numel` / `rnumel` become `_IntLike`, but both tiling dicts in the same
NamedTuple are deliberately left as they were, for two different reasons.

For `tiling_scores` the accurate value type is `int`
(`_SubSplit.split_scores: list[int]`, forced through `int(...)` at
`simd.py:5128`, zipped in by `create_tiling` without sympifying). `dict` is
invariant, so correcting it without also moving `SIMDKernel.__init__`,
`SIMDKernelFeatures`, `CandidateTiling`, `create_tiling`,
`ComboKernel.create_triton_kernel`, `FusedUserDefinedTritonKernel`, and
`BaseSchedulerNode.get_tiling` would leave that chain's declarations
contradicting each other -- a coherent follow-up, not this PR.

`tiling` is a tighter bind rather than the same one. Its producer and consumer
both already declare `dict[str, sympy.Expr]`, so matching them would cost
nothing and would make the third `_size_hint` call site checked -- but it would
not be true, since `create_tiling` also receives `tiling_factor` and the
`prod = 1` initializer, both plain `int`. The accurate `dict[str, _IntLike]` is
the one that collides with the chain. Since this PR's whole argument is that
these annotations should mean something, asserting an element type its own
evidence contradicts is the wrong trade, so `tiling` stays bare too.

Flagging the consequence rather than leaving it silent: `_size_hint`'s
narrowing is load-bearing as documentation at two of its three call sites, not
three.

Suggested reading order: `kernel/gemm_epilogue_utils.py` first, since it holds
the base predicates the rest call; then its callers in
`kernel/gemm_epilogue_analysis.py`, `kernel/flex_gemm/constraints.py`, and
`kernel/flex_gemm/lowering.py`; then the independent files (`kernel/mm.py`,
`heuristics/template/triton.py`, `optimize_indexing.py`, and
`codegen/triton_combo_kernel.py` together with the `codegen/simd.py` field
types it reads).

Three changes are not purely annotations.

`statically_known` used to fall through to
`symbolic_shapes.statically_known_true`, which accepts `bool | SymBool` and
raises `AssertionError` on anything else. With the parameter narrowed to
`object` that call no longer type checks. Rather than a `cast`, the fallthrough
now checks `isinstance(expr, torch.SymBool)` and raises `AssertionError` itself
-- the same outcome the callee already produced for the same inputs, but it
states the helper's real contract (`bool | sympy.Basic | SymBool`) in the code.
In practice the new branch is unreachable: every caller passes a comparison
result, which is a `bool` for concrete operands and for `sympy.Expr` `==`
(`Expr.__eq__` is structural), a `sympy` relational for `sympy.Expr` `<`/`<=`
and friends, and a `SymBool` for `SymInt` operands.

`grouped_tensor_layout` called `normalize_shape(shape[0])` inside a branch that
had already established `shape[0]` is a `list | tuple | torch.Size`, so
`normalize_shape` could only ever return `tuple(shape[0])` there. Calling
`tuple` directly keeps `shape` narrowed to a tuple for the rest of the function,
which is what the functions it is then handed to
(`_guard_grouped_reshape_group`, `_syntactic_grouped_tensor_layout`) already
declare they want.

`is_desired_scaling` and `get_scaling_options` in `kernel/mm.py` declared
`scale_size: torch.Tensor`, which was simply wrong -- every caller passes a size
(`scale_a_real.shape`, `scale_a.get_size()`) and the body indexes it and takes
its `len`. Those are now `Sequence[_IntLike]`.

One alternative considered and rejected: a single repo-wide alias for the
"maybe-symbolic scalar" union. These helpers do not agree on the union -- the
`isinstance`-branching ones genuinely accept `fx.Node`, the flex_gemm
local-reduce validators see only `int | SymInt`, and the rest see
`int | sympy.Expr` -- so one alias would have had to be the widest of the three
and would have re-hidden the distinction that motivated the change.

No `cast` and no new type-checker suppression were needed anywhere.

## Test Plan

No behavior change is intended, so the primary check is that the type checker
sees no new errors anywhere in the repo -- not just in the touched files, since
these are signature changes that propagate to callers:

```
pyrefly check   # with the patch, and again with the nine files reverted
```

Both runs report `65 errors (10,965 suppressed)`, and the sorted diagnostic
lists diff clean: zero net-new. Because the trailing count line is part of the
comparison, a newly *suppressed* error would also have shown up. Coverage of the
touched files was confirmed rather than assumed, by appending a deliberately
ill-typed function to `kernel/gemm_epilogue_utils.py` and checking the
whole-repo run reports it.

One caveat on how much that proves, since it cuts against this PR's own
motivation and is worth stating plainly. `pyrefly.toml:134` sets
`replace-imports-with-any = ["sympy.*", ...]`, so `sympy.Expr` resolves to `Any`
and `_IntLike` is effectively `int | Any`. I probed it directly -- a throwaway
`def f(x: _IntLike)` called with a `str` -- and it produces no diagnostic, while
the same probe against `IntLikeType` does. So the `_IntLike` annotations here
are accurate documentation that will begin checking if sympy stubs are ever
enabled, not checking today. The `object`, `IRNode`, `Sequence[object]`, and
`Sequence[IntLikeType]` narrowings are enforced now. `statically_known_multiple`'s
`_IntLike | IntLikeType` is the one exception on that side: the `_IntLike`
member absorbs everything, so the union documents rather than checks.

Lint, clean on all nine touched files:

```
lintrunner --take PYFMT,RUFF,PYREFLY,CODESPELL,FLAKE8 -a \
  torch/_inductor/codegen/simd.py \
  torch/_inductor/codegen/triton_combo_kernel.py \
  torch/_inductor/heuristics/template/triton.py \
  torch/_inductor/kernel/flex_gemm/constraints.py \
  torch/_inductor/kernel/flex_gemm/lowering.py \
  torch/_inductor/kernel/gemm_epilogue_analysis.py \
  torch/_inductor/kernel/gemm_epilogue_utils.py \
  torch/_inductor/kernel/mm.py \
  torch/_inductor/optimize_indexing.py
```

Tests, on a local build of this branch:

```
python test/inductor/test_optimize_indexing.py
python test/inductor/test_indexing.py
python test/inductor/test_flex_gemm.py
python test/inductor/test_torchinductor.py -k cpu
```

`test_optimize_indexing.py` is `OK` (15 run), `test_indexing.py` is `OK` (93
run, 22 skipped), and `test_flex_gemm.py` is `OK` (437 run, 369 GPU-skipped) --
the last includes `test_swap_ab_alignment_filters_tuned_and_explicit_configs`,
which is the test that pushes a `sympy.Symbol` through the corrected
`flex_gemm_config_keys` signature. The CPU inductor suite runs 1493 with 4
errors, all four `TritonMissing` from this box having no Triton install rather
than anything in the diff. The GPU-only suites that cover the rest of the touched code
(`test_flex_gemm.py`, `test_combo_kernels.py`, and the scaled-mm paths) are left
to CI.

This PR was authored with the assistance of an AI coding assistant.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/194482
Approved by: https://github.com/jansel