  This report presents some early results on code generation targeting tensor
cores on NVIDIA GPUs using the MLIR compiler infrastructure. The
state-of-the-art in high-performance deep learning today is primarily driven by
manually optimized highly tuned libraries. The approach to develop such
libraries is often not modular or reusable to the same extent that compiler
infrastructure like LLVM is. Manual optimization typically does not use a
standard intermediate representation (IR), although the optimizations performed
can be encoded as a sequence of transformation steps and customized passes on
an IR. Hand tuning may also miss exploration of design points only reachable
easily by automatic code generation. We believe that until the recent
introduction of MLIR (Multi-level intermediate representation), IR
infrastructure was not geared to tackle the problem of automatic generation of
domain-specific libraries in an effective manner. In particular, it was hard to
represent and transform compute abstractions at high, middle, and low levels
using a single IR.
  With suitable abstractions in MLIR, we build an experimental lowering
pipeline that is able to automatically generate code for matrix-matrix
multiplication on NVIDIA GPUs targeting its tensor cores. On a set of problem
sizes we evaluated, initial performance results show that we are able to attain
performance that is 95-119% and 80-160% of CuBLAS for FP32 and FP16 accumulate
respectively on NVIDIA's Ampere microarchitecture-based Geforce 3090 RTX. We
believe that these results could be used as motivation for further research and
development on automatic code and library generation using IR infrastructure
for similar specialized accelerators.
