Skip to content

Performancelink

Running models on RISC-V with IREE

IREE compiles machine learning models to native RISC-V CPU code, with support for the RISC-V Vector extension (RVV), hand-written microkernels, and data-tiling. This post walks through the full flow for a model: importing it from PyTorch, compiling it for a RISC-V target, running it, and benchmarking the result.

All commands below run under qemu-riscv64. The flow on real hardware is identical — the QEMU invocation is simply replaced by running the tools natively on the target.

Virtual Dense MFMAs for Skinny GEMM

When we have a GEMM A * B = C, and it is the situation that A has a small number of rows and many columns, we classify this problem as a skinny GEMM. The decode phase of LLM inference is a common sight of this problem: a small batch of tokens multiplies against a large weight matrix. Skinny GEMMs are less convenient for modern GPU architectures than their non-skinny cousins. One reason is because modern GPUs take advantage of matrix core units which offer instructions that are specifically designed for matrix multiplication and operate on fixed tile sizes, and skinny GEMMs are too small to utilize them to their intended size.

On AMDGPUs and in particular on the MI3XX Instinct (CDNA) series, these instructions are known as MFMA instructions; for example, V_MFMA_F32_16x16x16_F16. One useful part of the name is the MxNxK tile shape consumed, where M is the number of rows of the left hand matrix, N is the number of columns of the right hand matrix, and K is the shared dimension of both.

Data-Tiling Walkthrough

Data-tiling is the modification of data layout of operands of certain operations, such as matrix multiplication, that prefer specific layouts. These layout preferences depend on the operations and the target hardware. For example, matrix multiplications may need to use hardware matrix multiplication instructions that perform optimally with a specific matrix data layout.

Layout changes may also be motivated by memory access performance, as data-tiling can result in improved locality of memory accesses, fewer cache lines being accessed, and generally simpler memory access patterns that are more likely to be handled performantly by the memory system.

These layout changes can be propagated as far as possible across the workload, so the entire workload can use the updated layouts, as opposed to having to perform layout transformations at runtime. This may involve fusions or constant-evaluation that can amortize or remove layout-transformation overheads.

The main conceptual difficulty in modeling this in tensor-level MLIR is that tensors don't have layouts: tensors are higher-level, abstract arrays. This is addressed by the concept of tensor encodings, which this document will explain.

IREE / MLIR / Linalg tutorial

Introduction

This tutorial is simultaneously about IREE, MLIR, and specifically the MLIR Linalg dialect.

What is MLIR?

MLIR is a programming language, but MLIR in itself is almost just an empty shell. What it really provides is a framework allowing to define MLIR dialects which are where the features come from.

The "IR" part of the MLIR name stands for "intermediate representation". It means that MLIR is meant to be primarily for compiler-internal representations of code. But MLIR is actually fairly nice for humans to work with, and it's not hard to hand-author some MLIR programs from scratch. That is exactly the topic of this tutorial.

Exploring CPU microkernels on a matmul example

Basic setup, command lines

Source file: matmul.mlir:

func.func @matmul_dynamic(%lhs: tensor<?x?xf32>, %rhs: tensor<?x?xf32>, %acc: tensor<?x?xf32>) -> tensor<?x?xf32> {
  %result = linalg.matmul ins(%lhs, %rhs: tensor<?x?xf32>, tensor<?x?xf32>) outs(%acc: tensor<?x?xf32>) -> tensor<?x?xf32>
  return %result: tensor<?x?xf32>
}

Basic compilation command line:

$ iree-compile matmul.mlir -o /tmp/matmul.vmfb \
  --iree-hal-target-device=local \
  --iree-hal-local-target-device-backends=llvm-cpu \
  --iree-llvmcpu-target-cpu=znver4 \
  --iree-llvmcpu-enable-ukernels=all

This creates a IREE bytecode module:

$ ls -l /tmp/matmul.vmfb

-rw-rw-r-- 1 2884 Jan 22 10:37 /tmp/matmul.vmfb

Matrix Multiplication with MMT4D

Introduction

Matrix multiplication (matmul) is an important operation in ML workloads that poses specific challenges to code generation. For example, matmul makes repeated accesses to the same data, which makes locality of reference a top concern.

Moreover, modern CPUs instruction set architectures (ISAs) offer specialized SIMD instructions that the matmul implementation needs to use to achieve optimal performance, and these instructions expect data to be in a particular layout.

This article is about an in-development MLIR operation, linalg.mmt4d, offering a compilation path for linalg.matmul that is designed from the ground up for these efficiency considerations.