arXiv subjects

Justs Zarins

Publications and source records attributed to Justs Zarins.

3 recordsLinked to original sources

arXiv2026

An MLIR Lowering Pipeline for Stencils at Wafer-Scale

The Cerebras Wafer-Scale Engine (WSE) delivers performance at an unprecedented scale of over 900,000 compute units, all connected via a single-wafer on-chip interconnect. Initially designed for AI, the WSE architecture is also well-suited for High Performance Computing (HPC). However, its distributed asynchronous programming model diverges significantly from the simple sequential or bulk-synchronous programs that one would typically derive for a given mathematical program description. Targeting the WSE requires a bespoke re-implementation when porting existing code. The absence of WSE support in compilers such as MLIR, meant that there was little hope for automating this process. Stencils are ubiquitous in HPC, and in this paper we explore the hypothesis that domain specific information about stencils can be leveraged by the compiler to automatically target the WSE without requiring application-level code changes. We present a compiler pipeline that transforms stencil-based kernels into highly optimized CSL code for the WSE, bridging the semantic gap between the mathematical representation of the problem and the WSE's asynchronous execution model. Based upon five benchmarks across three HPC programming technologies, running on both the Cerebras WSE2 and WSE3, our approach delivers comparable, if not slightly better, performance than manually optimized code. Furthermore, without requiring any application level code changes, performance on the WSE3 is around 14 times faster than 128 Nvidia A100 GPUs and 20 times faster than 128 nodes of a CPU-based Cray-EX supercomputer when using our approach.

cs.DC↗

arXiv2024

Pragma driven shared memory parallelism in Zig by supporting OpenMP loop directives

The Zig programming language, which is designed to provide performance and safety as first class concerns, has become popular in recent years. Given that Zig is built upon LLVM, and-so enjoys many of the benefits provided by the ecosystem, including access to a rich set of backends, Zig has significant potential for high performance workloads. However, it is yet to gain acceptance in HPC and one of the reasons for this is that support for the pragma driven shared memory parallelism is missing. In this paper we describe enhancing the Zig compiler to add support for OpenMP loop directives. Then exploring performance using NASA's NAS Parallel Benchmark (NPB) suite. We demonstrate that not only does our integration of OpenMP with Zig scale comparatively to Fortran and C reference implementations of NPB, but furthermore Zig provides up to a 1.25 times performance increase compared to Fortran.

cs.DC↗

arXiv2022

TensorFlow as a DSL for stencil-based computation on the Cerebras Wafer Scale Engine

The Cerebras Wafer Scale Engine (WSE) is an accelerator that combines hundreds of thousands of AI-cores onto a single chip. Whilst this technology has been designed for machine learning workloads, the significant amount of available raw compute means that it is also a very interesting potential target for accelerating traditional HPC computational codes. Many of these algorithms are stencil-based, where update operations involve contributions from neighbouring elements, and in this paper we explore the suitability of this technology for such codes from the perspective of an early adopter of the technology, compared to CPUs and GPUs. Using TensorFlow as the interface, we explore the performance and demonstrate that, whilst there is still work to be done around exposing the programming interface to users, performance of the WSE is impressive as it out performs four V100 GPUs by two and a half times and two Intel Xeon Platinum CPUs by around 114 times in our experiments. There is significant potential therefore for this technology to play an important role in accelerating HPC codes on future exascale supercomputers.

cs.DC↗