Extending LLVM IR for DPC++ Matrix Support: A Case Study with Intel® Advanced Matrix Extensions (Intel® AMX)

In this paper, we introduce a DPC++ matrix extension to unify different tensor hardware: Intel
^®
Advanced Matrix Extensions (Intel
^®
AMX) to CPUs, NVIDIA
^®
TPUs, IBM
^®
POWER
^®
MMA, etc. These tensor hardware units are usually accessed by low-level intrinsics or assembly to perform matrix operations. It is hard for scientists to program these domain- specific devices without the kind of high-level abstractions and efficient implementations we introduce here.We also extend the existing LLVM matrix intrinsics to represent this DPC++ extension and yield efficient Intel AMX code generation. Based on our case study of implementing this interface on Intel AMX hardware, we discuss some of the limitations of existing LLVM Intermediate Representation (IR) and how they can be overcome to exploit tensor hardware.