Adnan Akhundov 06d2392003 Support tt.reduce in Triton kernel analysis pass (#121706) ... Summary: Previously, we bailed out of the Triton kernel analysis pass when seeing a `tt.reduce` op. In this PR, we support the op and don't bail out anymore. Test Plan: This is a bit tricky, as the extension is added to the MLIR walk-based analysis code path which is active only on when the MLIR bindings added in https://github.com/openai/triton/pull/3191 are available. So for now I've run the `test_argmax` and `test_reduce_sum` manually with a newer Triton version than the current pin. When pin updates, we'll make those tests official (left a TODO comment). Pull Request resolved: https://github.com/pytorch/pytorch/pull/121706 Approved by: https://github.com/jansel		2024-03-12 11:38:28 +00:00
..
__init__.py	[HigherOrderOp] expose torch.cond (#110293)	2023-10-07 20:39:52 +00:00
auto_functionalize.py	Support higher order op functionalization in predispatch IR (#115314)	2024-03-01 09:13:47 +00:00
cond.py	Support higher order op functionalization in predispatch IR (#115314)	2024-03-01 09:13:47 +00:00
effects.py	[export] Add effect token to export (#121424)	2024-03-09 02:43:26 +00:00
map.py	Clean up mode handling in python dispatcher (#121083)	2024-03-08 00:30:34 +00:00
out_dtype.py	Support higher order op functionalization in predispatch IR (#115314)	2024-03-01 09:13:47 +00:00
strict_mode.py	[torch.export] Support is_compiling() flag for non-strict mode (#119602)	2024-02-29 05:52:51 +00:00
torchbind.py	[export] improve FakeTensor serialization (#119531)	2024-02-12 19:28:08 +00:00
triton_kernel_wrap.py	Support tt.reduce in Triton kernel analysis pass (#121706)	2024-03-12 11:38:28 +00:00
utils.py	Support higher order op functionalization in predispatch IR (#115314)	2024-03-01 09:13:47 +00:00
while_loop.py	Support higher order op functionalization in predispatch IR (#115314)	2024-03-01 09:13:47 +00:00
wrap.py	[export] add replace_set_grad_with_hop_pass (#119810)	2024-02-17 02:18:19 +00:00