iree-org / iree-org/iree

[MLIR] [uKernel] Implement MLIR-based Ukernels for RISC-V Vector (RVV) Extension

Open
#22,720 3 comments 0 reactions 0 assignees View on GitHub
enhancement ➕
Dominant language
C++
Stars
3.9k
Forks
1k
Avg merge
4d 16h
Merged PRs (30d)
47

Description

### Request description

## Objective

Implement MLIR-native ukernels for RISC-V RVV targets in IREE. This follows the GPU backend's precedent of using declarative, tensor-based ukernels. This design enables high-performance, target-specific code generation while preserving IREE's ability to perform automatic producer/consumer fusion (e.g., bias addition and activation) via standard compiler passes.

I plan to rely on the generic vector dialect for the implementation, rather than introducing a target-specific dialect, delegating instruction selection to the LLVM backend.

## MVP Scope

The Minimum Viable Product will be a single, pure `matmul` ukernel for the `f16 * f16 -> f32` mixed-precision case, targeting a `vlen=128` RVV configuration with the `zfh` extension. This single ukernel will be used to demonstrate automatic fusion for a `matmul+bias` pattern.

## Proposed Engineering Plan

### 1. Ensure Robust Lowering from `vector` Dialect to RVV in LLVM

We should rely on the `vector` dialect. The responsibility for generating optimal RVV instructions lies with the LLVM backend.

* **Action:** Verify that `vector` dialect operations are correctly lowered to efficient RVV instructions by the MLIR-to-LLVM-IR conversion and LLVM's instruction selector. This includes:
* `vector.fma %a, %b, %c` -> `vfmacc.vf` (for vector-scalar)
* `vector.load`, `vector.store` -> `vle16.v`, `vse32.v`
* `vector.broadcast` -> `vmv.v.x`
* **Upstream Contribution:** If gaps are found, the work involves improving the lowering patterns in upstream MLIR/LLVM, not creating a new dialect in IREE.

### 2. Implement IREE Infrastructure for CPU Ukernel Declaration

To match the GPU backend's capabilities, I need to add declarative attributes for CPU targets:

1. **Define a CPU Layout Attribute:** Create an `#iree_cpu.data_tiled_layout` MLIR attribute to describe the optimal data tiling for a CPU ukernel. This tells IREE's `MaterializeDeviceEncoding` pass how to prepare the data.
```mlir
// Example for our MVP ukernel
#rvv_mvp_layout = #iree_cpu.data_tiled_layout<
tile_sizes = [8, 8, 1], // M0=8, N0=8, K0=1
inner_blocking_order = [0, 1, 2]
>
```
2. **Define a CPU Feature Matcher:** Implement a `cpu_feature` constraint for the ukernel matching system. This allows a ukernel to be selected only when the target CPU has the required features (e.g., `+v`, `+zfh`).

### 3. Define and Implement the RVV Matmul Ukernel in MLIR

It is a self-contained `.mlir` file that is both the specification and implementation of the kernel.

* Create the `iree_uk_riscv_dt_matmul_f16f16f32.mlir` file with the following structure:

```mlir
// Ukernel for f16*f16->f32 on RVV. Bias is handled by compiler fusion.

// --- Declarative Metadata ---
#rvv_f16_layout = #iree_cpu.data_tiled_layout<
tile_sizes = [8, 8, 1], inner_blocking_order = [0, 1, 2]>

iree_codegen.ukernel.generic @riscv_matmul_8x8x1_f16f16f32
match="linalg.matmul"
with
operand_types, tensor<...xf16>, tensor<...xf32>> and
cpu_feature<"+v", "+zfh">
provides layout<#rvv_f16_layout>
{
// --- Ukernel Implementation (using generic vector ops) ---
util.func public @main(
%lhs: tensor<8x1xf16>, %rhs: tensor<1x8xf16>, %init: tensor<8x8xf32>) -> tensor<8x8xf32>
{
%vtype_f16 = vector<8xf16>
%vtype_f32 = vector<8xf32>

// Load RHS vector (reused for all 8 rows of LHS).
%rhs_vec = vector.load %vtype_f16, %rhs[0, 0] : tensor<1x8xf16>, vector<8xf16>

// FUSION POINT: Load initial accumulators. The compiler will feed the
// bias tensor into the '%init' argument.
%acc0 = vector.load %vtype_f32, %init[0, 0] : tensor<8x8xf32>, vector<8xf32>
%acc1 = vector.load %vtype_f32, %init[1, 0] : tensor<8x8xf32>, vector<8xf32>
%acc2 = vector.load %vtype_f32, %init[2, 0] : tensor<8x8xf32>, vector<8xf32>
%acc3 = vector.load %vtype_f32, %init[3, 0] : tensor<8x8xf32>, vector<8xf32>
%acc4 = vector.load %vtype_f32, %init[4, 0] : tensor<8x8xf32>, vector<8xf32>
%acc5 = vector.load %vtype_f32, %init[5, 0] : tensor<8x8xf32>, vector<8xf32>
%acc6 = vector.load %vtype_f32, %init[6, 0] : tensor<8x8xf32>, vector<8xf32>
%acc7 = vector.load %vtype_f32, %init[7, 0] : tensor<8x8xf32>, vector<8xf32>

// Core matmul logic using generic vector FMA.
%lhs_s0 = memref.load %lhs[0, 0] : tensor<8x1xf16>
%b_lhs_s0 = vector.broadcast %lhs_s0 : f16 to vector<8xf16>
%res0 = vector.fma %b_lhs_s0, %rhs_vec, %acc0 : vector<8xf32>

%lhs_s1 = memref.load %lhs[1, 0] : tensor<8x1xf16>
%b_lhs_s1 = vector.broadcast %lhs_s1 : f16 to vector<8xf16>
%res1 = vector.fma %b_lhs_s1, %rhs_vec, %acc1 : vector<8xf32>

// ... (repeated for res2 through res7) ...

%lhs_s7 = memref.load %lhs[7, 0] : tensor<8x1xf16>
%b_lhs_s7 = vector.broadcast %lhs_s7 : f16 to vector<8xf16>
%res7 = vector.fma %b_lhs_s7, %rhs_vec, %acc7 : vector<8xf32>

// Assemble final result tensor.
%out0 = tensor.insert %res0 into %init[0, 0] : tensor<8x8xf32>
%out1 = tensor.insert %res1 into %out0[1, 0] : tensor<8x8xf32>
// ... (repeated for out2 through out7) ...
%out7 = tensor.insert %res7 into %out6[7, 0] : tensor<8x8xf32>

util.return %out7 : tensor<8x8xf32>
}
}
```

### 4. Integrate Ukernel into IREE's Build and Codegen Pipeline

1. Place the ukernel `.mlir` file in a directory like `iree/compiler/Codegen/LLVMCPU/builtins/mlir_ukernel/riscv/`.
2. Add a PDL pattern in `iree/compiler/Codegen/LLVMCPU/builtins/pdl/` to match the `linalg.matmul` and rewrite it using the `iree_codegen.match.ukernel` rewriter.
3. Ensure the `LowerTensorUKernelsPass` is part of the LLVMCPU codegen pipeline and correctly finds and inlines the ukernel's `@main` function.

### 5. Define a Comprehensive Test Suite

1. **Integration Test:** Add an IREE test case that compiles a `linalg.matmul` with a preceding `linalg.fill` (for bias). The test should target RVV and check the IR after the `LowerTensorUKernelsPass` to verify that the ukernel was inlined and that the `linalg.fill` is correctly feeding its output into the ukernel's `init` argument.
```mlir
// test/iree/compiler/Codegen/LLVMCPU/riscv_matmul_bias_fusion.mlir
// RUN: iree-compile ... --iree-llvmcpu-target-triple=riscv64... --iree-llvmcpu-target-cpu-features="+v,+zfh" ...
// CHECK: func.func private @riscv_matmul_8x8x1_f16f16f32_dispatch
// CHECK: vector.load %{{.*}} : memref<...>, vector<8xf32>

func.func @test_matmul_bias(%lhs: tensor<512x256xf16>, %rhs: tensor<256x1024xf16>, %bias: tensor<1024xf32>) -> tensor<512x1024xf32> {
%init = linalg.fill ins(%bias) outs(...) : tensor<512x1024xf32>
%res = linalg.matmul ins(%lhs, %rhs) outs(%init) -> tensor<512x1024xf32>
return %res : tensor<512x1024xf32>
}
```
2. **End-to-End Test:** Create an e2e test that runs the compiled module on an emulator (e.g., QEMU) or hardware and verifies the numerical correctness of the fused `matmul+bias` operation against a reference implementation.

### What component(s) does this issue relate to?

MLIR

### Additional context

_No response_

Contributor guide

Open the contributing guide

Research direction

Start by reading the LLVMCPU builtins paths under iree/compiler/Codegen/LLVMCPU/builtins/mlir_ukernel/riscv/ and builtins/pdl/, then inspect the LowerTensorUKernelsPass. Add the RVV ukernel, CPU layout and feature matching, pipeline integration, and the proposed riscv_matmul_bias_fusion.mlir test. Done requires the fusion integration check and an end-to-end numerical test on an emulator or hardware.

Written by the indexing model from the issue text.

Assessment

Domain
compilers
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.