Example Conv2D in Triton

[sebastienwood]

Hi !
I'm interested in implementing a Conv-Nd-like operation in Triton. I'm pretty new to GPU programming and Triton. The associated paper proposed a Conv2d implementation in C-Triton. Is there a Python equivalent available ?
My main technical questions lies in the computation of the offsets and the general optimization strategy. Should there be two or more program ids ? (I'm not so sure what is the use of tl.program_id and the documentation is not so clear on it).

My end goal is not amenable to unfolding. To be precise, I only need to compute the following equation (+ batched):
$Z_{a, b, c, a, d, e} = \sum_g \sum_i \sum_j \gamma_{g, b+i, c+j, g, d+i, e+j}$
Let me know if I can provide further information !

[yangjunpro]

I think you can try to start from here
We have tried it locally and functionally it works with ResNet-50 inference problem size.
However, the performance is not good(we have seen around 2X slow-down compared against the PyTorch native implementation which calls cuDNN instead). Of course, for popular workload like RN50, we should not expect significant perf gain with Triton, but the 2X perf slow-down still deserves some investigation, which is our current ongoing work.

BTW, as to conv performance, do you have any previous empirical data to be shared? @ptillet

[ptillet]

Yeah, I remember investigating this code, and finding out that the issue came from lack of vectorization, which snowballed and made async copies inapplicable. Will be worth re-investigating once the new MLIR backend lands.

[sebastienwood]

Thanks for the detailed answers ! Out of curiosity:

• do you have any ETA for the new MLRIR backend ?
• I understand MLIR is a new layer to the LLVM ecosystem, and is taylor made for ML/DL. How does it impact Triton ? Does it makes some computation graphs easier to optimize ? Do I miss some other "cool stuff" behind the scene ?
• Is there some readings you would recommand to get into understanding/contributing to Triton ? (LLVM's tutorial seems like the first step in my naive eyes)

[ptillet]

• I don't want to get ahead of myself, but it's making good progress and we're hoping it'll land in ~2months.
• I wouldn't say MLIR is taylor made for DL or even Linear Algebra. It's really just nice infrastructure for representing and transforming programs at different levels of abstractions. We don't use it for anything fancy; only to avoid reinventing the wheel on traditional compiler infra components.
• I think I would just recommend to read the Triton-MLIR codebase (as well as the MLIR docs) and follow the pull requests that are being made there :)

[jsonlee0x02]

Seems some efforts have been done in the current version of Triton.
But still cannot find any convolution example in the codebase of latest triton?
BTW, any updates to the performance of ResNet50?

[flishwang]

Two years have been passed, are there any updates?

[l1351868270]

@sebastienwood @ptillet you see "The associated paper proposed a Conv2d implementation in C-Triton." i implement one in python, but it performance is poor! i do not know how to improve it, could you tell me which paper。i want to refer to it to improve some performance https://github.com/l1351868270/implicit_gemm.triton/blob/main/triton_implicit_gemm.py

[l1351868270]

i found why the performance is poor? when the data load from global memory to shared memory, the ptx code do not use the cp.async future. could i force the tl.load compile to use cp.async?

ps:
i write a test code with U=1, V=1, pad_h=0, pad_w=0, dila_h=1, dila_w=1, then the conv2d is same as the gemm only the shape is different. When i use x_ptrs, w_ptrs the performance is poor, and when i use a_ptrs and b_ptrs the performace is ok.
the source code is here https://github.com/l1351868270/implicit_gemm.triton/blob/main/triton_implicit_gemm_1x1_0x0_1x1.py

command:
python python triton_implicit_gemm_1x1_0x0_1x1.py
python triton_bench_1x1_0x0_1x1.py
ncu -f --set full --call-stack -o bench_conv2d_report python triton_implicit_gemm_1x1_0x0_1x1.py

[l1351868270]

i write a new conv2d use triton that support forward and backward https://github.com/l1351868270/ld_triton/blob/main/ld_triton/ops/convolution/triton_conv2d.py

[webstorms]

Thanks for sharing! Are there any benchmarks to see how this compares to PyTorch?