TinyTC and PyTorch/Triton boost stencil computing on GPUs
Stencil Computation at the Intersection of AI and HPC
Distributed, Parallel, and Cluster Computing
Summary
Stencil computations are a type of calculation important for science and engineering, often running on powerful GPUs. The authors show that tools originally made for AI, like TinyTC and PyTorch/Triton, can also speed up these stencil calculations. They tested these tools on many different hardware setups, finding that TinyTC is fastest on some Intel chips, while PyTorch/Triton works well across many types of GPUs. They also found that the main limit on speed is how fast data can move in memory, not the calculation speed itself. This helps people choose the best tools for running complex simulations efficiently.
What this means in practice
- •For high performance computing teams: Optimize complex stencil-based scientific simulations by using TinyTC for better performance on Intel GPUs or PyTorch/Triton for more portable implementations.
- •For industrial simulation engineers: Develop efficient acoustic or wave propagation models on diverse GPU hardware using AI tensor compilers originally designed for neural networks.
Authors
Timothee Ewart, Mauricio Araya-Polo
Abstract
Tensor compilers such as TinyTC and OpenAI Triton were originally developed for AI workloads, but the same tiling and memory abstractions can be applied to implement efficient high-order stencils for scientific and industrial applications. We demonstrate this for an 8th-order, 25-point acoustic stencil with boundary conditions over an a demanding-sized grid, targeting GPGPUs, where we compare the hardware-specialized TinyTC implementation with a portable PyTorch/Triton implementation. The target platforms for evaluation include Intel B70, B580, GPU MAX 1550, NVIDIA A100/RTX6000 Blackwell/H100, and AMD MI325x. For instance, on Battlemage B580 TinyTC reaches 15.6 Gpts/s versus 13.5 Gpts/s for PT/Triton under random initialization, while zero-initialized runs reach up to 35.8 Gpts/s due to hardware memory compression. Using roofline and memory-hierarchy profiling, we show that -as expected- performance is predominantly bandwidth-limited and that compiler-managed L1/LSC caching can effectively replace programmer-managed shared-memory staging for this stencil class. Overall, the results position TinyTC as the performance-oriented path on Intel hardware and PyTorch/Triton as a strong portability/productivity baseline for cross-vendor HPC stencil development.