trunk/9c6b1aa07d56296da76db2b6c429d7b32efbe2eb: Triton backward convolution kernel (#178945)

This PR implements Triton kernels for 2D convolution backward operations (input and weight gradients) in TorchInductor, replacing the previous ATen-only fallback approach.

Key Changes

Core Implementation:

• Added Triton templates for conv2d_bwd_weight and conv2d_bwd_input operations


• Implemented convolution_backward_lowering function with backend selection logic


• Added layout computation functions for backward operations


• Added configuration options for backend selection (max_autotune_conv_bwd_weight_backends, max_autotune_conv_bwd_input_backends)

Test Coverage:

• Added comprehensive test_conv2d_backward_parametrized test covering various parameter combinations (stride, dilation, padding, kernel sizes, groups, NHWC format)


• Includes proper tolerance handling for ROCM vs other platforms


• Added skip conditions for unsupported dynamic shapes and subprocess compilation mode


• Removed test_conv2d_backward_channels_last_dynamic_shapes from test failures (now passing)

Performance Impact

• For certain models, small-batch workloads can achieve up to 20% end-to-end performance improvement on MI325 GPU. Of the 20% performance increase, about 3% comes from reduced kernel execution time, and 17% comes from reduced bubble time.


• For a single operator, the triton bwd conv kernel has a good chance of outperforming the aten implementation in scenarios input_H=1 input_W=1 kernel_H=1 kernel_W=1 for backward weight, kernel_H=1 kernel_W=1 for backward input. Below are some shape/params that the triton bwd conv kernel has better performance than aten. The data type is bf16, and the tests is run on MI325 GPU.

Pull Request resolved: , https://github.com/eellison

Co-authored-by: Eric.Chin.AMD [email protected]