SIGN IN SIGN UP

[mtia][inductor] resolve view args of user-defined Triton kernels in FX wrapper (#193820) (#193820)

Summary:

Compiling a model that calls a hand-written Triton kernel fails with
`InductorError: KeyError: 'reinterpret_tensor(buf0, (8,), (1,), 0)'` whenever
one of the kernel's tensor arguments is a view of another buffer, such as a
slice of a matmul result. It happens on the FX wrapper backend, the one that
emits an FX graph instead of Python source: that backend could only map a
plain buffer name back to a graph node, so a view argument, which code
generation renders as a `reinterpret_tensor(...)` call string, matched nothing
and the compile aborted. The kernel call already carries the IR node behind
every argument, so this builds the view's node from that whenever the name
lookup misses. Plain buffer names take exactly the path they did before.

Test Plan:
The change is in the FX wrapper backend's argument lookup, so both tests feed a
user-defined Triton kernel two halves of a matmul result, which reach code
generation as the `reinterpret_tensor(...)` views this diff resolves. The OSS
test compiles such a model through that backend end to end; the MTIA one runs
the same shape through the full inductor lowering. Both fail with this diff
reverted, raising the `KeyError` from the summary. The bug came from production
model `879326229` (snapshot 249), which now compiles past code generation.

```bash
buck2 test fbcode//mode/opt fbcode//caffe2/test/inductor:fxir_backend \
  -- --regex 'test_custom_triton_view_arg'
```

Buck UI: https://www.internalfb.com/buck2/da7696ca-fd98-4c69-ac71-9c440cff5a96
Infra: https://www.internalfb.com/intern/testinfra/testrun/844425452630609

```bash
buck2 test fbcode//mode/opt -m ovr_config//gpu:mtia \
  fbcode//glow/fb/fx/fba/tests/inductor:test_fba_inductor \
  -- --regex 'test_full_inductor_user_triton_kernel_view_arg'
```

Buck UI: https://www.internalfb.com/buck2/a1f7c6bf-21e6-47d7-b3c9-99241a934887
Infra: https://www.internalfb.com/intern/testinfra/testrun/26458647852100141

Reviewed By: blaine-rister

Differential Revision: D116270888

Pull Request resolved: https://github.com/pytorch/pytorch/pull/193820
Approved by: https://github.com/blaine-rister
Z
Zhitao Lin committed
a6531c1ff713aeacb1936bae1d5e24989d478ca0
Parent: 0184680
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> on 8/25/2026, 5:43:30 PM