[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