mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-15 21:00:47 +00:00
b1b3f61f2c
Summary:
MM max autotune (and friends) crash when one of the inputs is zero-size.
E.g., running this code:
```
@torch.compile()
def fn(x, y):
return torch.mm(x, y)
inps = [torch.rand([0, 30]), torch.rand([30, 40])]
inps = [x.to(device="cuda") for x in inps]
out = fn(*inps)
```
with this command:
```
TORCHINDUCTOR_MAX_AUTOTUNE=1 python test.py
```
raises this error (the top of the stack trace omitted for brevity):
```
...
File "/data/users/aakhundov/pytorch/torch/_inductor/kernel/mm.py", line 119, in tuned_mm
return autotune_select_algorithm("mm", choices, [mat1, mat2], layout)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/data/users/aakhundov/pytorch/torch/_inductor/select_algorithm.py", line 960, in autotune_select_algorithm
return _ALGORITHM_SELECTOR_CACHE(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/data/users/aakhundov/pytorch/torch/_inductor/select_algorithm.py", line 787, in __call__
timings = self.lookup(
^^^^^^^^^^^^
File "/data/users/aakhundov/pytorch/torch/_inductor/codecache.py", line 267, in lookup
timings[choice] = benchmark(choice)
^^^^^^^^^^^^^^^^^
File "/data/users/aakhundov/pytorch/torch/_inductor/select_algorithm.py", line 774, in autotune
raise ErrorFromChoice(msg, choice, benchmark_fn.debug_str())
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
LoweringException: ErrorFromChoice: Please run `ptxas /tmp/compile-ptx-src-bfb1c6` to confirm that this is a bug in `ptxas`
From choice TritonTemplateCaller(/tmp/torchinductor_aakhundov/z7/cz7n7nn6rdlaelu4pbaaurgmu74ikl6g76lkngwawrevlfxlc6re.py, ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=16, BLOCK_N=64, EVEN_K=False, GROUP_M=8, num_stages=2, num_warps=4)
inputs = [
torch.empty_strided((0, 30), (30, 1), dtype=torch.float32, device='cuda'),
torch.empty_strided((30, 40), (40, 1), dtype=torch.float32, device='cuda'),
]
out = torch.empty_strided((0, 40), (40, 1), dtype=torch.float32, device='cuda')
target: aten.mm.default
args[0]: TensorBox(StorageBox(
InputBuffer(name='arg1_1', layout=FixedLayout('cuda', torch.float32, size=[0, s0], stride=[s0, 1]))
))
args[1]: TensorBox(StorageBox(
InputBuffer(name='arg3_1', layout=FixedLayout('cuda', torch.float32, size=[s0, s1], stride=[s1, 1]))
))
```
This PR adds a check to skip Triton templates in the `mm`, `addmm`, `mm_plus_mm` autotuning when the product of the MM problem shape (`m * n * k`) is zero.
Additionally, early exit conditions have been added to the mm and mm_plus_mm Triton templates on `M * N * K == 0`, to prevent issues when autotuning was done on non-zero-size inputs with dynamic shapes, then zero-size inputs are encountered by the compiled model.
Test Plan:
```
$ python test/inductor/test_max_autotune.py -v
...
----------------------------------------------------------------------
Ran 16 tests in 29.569s
OK
```
Reviewers: @eellison
Subscribers:
Tasks:
Tags:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/106865
Approved by: https://github.com/jansel
|
||
|---|---|---|
| .. | ||
| codegen | ||
| fx_passes | ||
| kernel | ||
| __init__.py | ||
| autotune_process.py | ||
| bounds.py | ||
| codecache.py | ||
| compile_fx.py | ||
| config.py | ||
| coordinate_descent_tuner.py | ||
| cuda_properties.py | ||
| cudagraph_trees.py | ||
| debug.py | ||
| decomposition.py | ||
| dependencies.py | ||
| exc.py | ||
| freezing.py | ||
| fx_utils.py | ||
| graph.py | ||
| hooks.py | ||
| index_propagation.py | ||
| inductor_prims.py | ||
| ir.py | ||
| lowering.py | ||
| metrics.py | ||
| optimize_indexing.py | ||
| pattern_matcher.py | ||
| scheduler.py | ||
| select_algorithm.py | ||
| sizevars.py | ||
| test_operators.py | ||
| triton_helpers.py | ||
| triton_heuristics.py | ||
| utils.py | ||
| virtualized.py | ||
| wrapper_benchmark.py | ||