# qmatmul

**Repository Path**: Nicet/qmatmul

## Basic Information

- **Project Name**: qmatmul
- **Description**: 四元数张量矩阵乘法
- **Primary Language**: Unknown
- **License**: Apache-2.0
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-07-07
- **Last Updated**: 2026-07-11

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# QMatmul 算子 — 四元数矩阵乘法 (Ascend 310B 版)

> 合作方：大湾区大学罗除课题组

## 算子介绍

本算子实现**四元数张量矩阵乘法**（Quaternion Matrix Multiplication），输入两个四元数张量，按四元数乘法规则计算得到输出张量。

| 张量 | 形状 | 说明 |
|------|------|------|
| A | `(M, K, 4)` | 左乘四元数矩阵 |
| B | `(K, N, 4)` | 右乘四元数矩阵 |
| C | `(M, N, 4)` | 输出四元数矩阵 |

最后一个维度 4 分别对应四元数的四个分量：实部 r、虚部 i、虚部 j、虚部 k。

---

## 背景与意义

### 四元数简介

四元数（Quaternion）是复数的推广，形式为 `q = r + xi + yj + zk`，其中 `i² = j² = k² = ijk = -1`。四元数在以下领域有广泛应用：

- **3D 旋转与姿态表示**：比欧拉角无死锁，比旋转矩阵更紧凑
- **计算机图形学**：骨骼动画、相机旋转
- **机器人学 / 航天**：IMU 姿态解算、飞行器控制
- **四元数神经网络**：将传统实数值网络扩展到四元数域，在彩色图像处理、语音情感识别等领域展现更好的表示能力

四元数矩阵乘法是上述应用中最核心、计算最密集的算子，在 NPU 上高效实现它对整体性能至关重要。

### 为什么需要自定义算子

四元数矩阵乘法**无法**用 Ascend 内置的矩阵乘算子直接表达。一方面，其乘法规则涉及各分量之间的交叉混合（Hamilton 乘积），而非简单的实数矩阵乘；另一方面，直接实现需要 16 次小矩阵乘法 + 12 次加法，计算量和带宽开销巨大。本算子通过巧妙的**矩阵展开（Expand）策略**，将四元数乘法转化为**单次标准实数矩阵乘法**，最大化利用 NPU Cube 单元的算力。

---

## 算法原理

### Hamilton 乘积规则

两个四元数 `Qa = r₁ + x₁i + y₁j + z₁k` 与 `Qb = r₂ + x₂i + y₂j + z₂k` 的乘积为：

```
C_r = r₁·r₂ - x₁·x₂ - y₁·y₂ - z₁·z₂
C_i = r₁·x₂ + x₁·r₂ + y₁·z₂ - z₁·y₂
C_j = r₁·y₂ - x₁·z₂ + y₁·r₂ + z₁·x₂
C_k = r₁·z₂ + x₁·y₂ - y₁·x₂ + z₁·r₂
```

对于 `(M, K, 4)` 和 `(K, N, 4)` 的矩阵乘法，naive 实现需要 16 次 `(M x K) @ (K x N)` 的小矩阵乘和 12 次加减法。

### 矩阵展开策略（核心思想）

本算子通过以下变换将四元数矩阵乘法降维为单次标准实数矩阵乘法：

**Step 1 — 零拷贝重塑 A**

A 在内存中按 `[r, i, j, k]` 交织排列，形状 `(M, K, 4)` 直接被视图为 `(M, 4K)` 的二维矩阵。**不发生数据搬运**，仅是内存视角的转换。

**Step 2 — 构造 B 的展开矩阵**

将 B 从 `(K, N, 4)` 展开为 `(4K, 4N)`。展开规则由 Hamilton 乘积决定，对应一个 4×4 分块矩阵映射：

| 展开行块 | r 列块 | i 列块 | j 列块 | k 列块 |
|:--------:|:------:|:------:|:------:|:------:|
| **r 行块** |  +B_r |  +B_i |  +B_j |  +B_k |
| **i 行块** |  -B_i |  +B_r |  -B_k |  +B_j |
| **j 行块** |  -B_j |  +B_k |  +B_r |  -B_i |
| **k 行块** |  -B_k |  -B_j |  +B_i |  +B_r |

即展开矩阵中 `[ii*K + i, jj*N + j]` 处的值由 `B[i, j]` 的第 `idx_map[ii][jj]` 个分量乘以 `sign_map[ii][jj]` 算出。此过程在 NPU 的 Vector 单元上通过循环 + 查表完成。

**Step 3 — 单次 Cube 矩阵乘**

在 Cube 单元上执行一次标准的实数矩阵乘法：
```
C_flat = A_flat @ B_expand     // shape: (M, 4K) @ (4K, 4N) → (M, 4N)
```

**Step 4 — 零拷贝输出**

将结果 `(M, 4N)` 直接视图为 `(M, N, 4)`，无需数据重排。

> **验证**：该算法在 Python 中经过 `torch.allclose` 验证，与 naive 四元数乘法的结果完全一致。

---

## 实现方案（Ascend 310B 版本）

### 整体架构

```
┌─────────────────────────────────────────────────────────────────┐
│  用户调用层 (aclnnQMatmul / AclNNInvocationNaive)               │
├─────────────────────────────────────────────────────────────────┤
│  op_host/q_matmul.cpp          — 算子注册、原型定义、shape推导   │
│  op_host/q_matmul_tiling.h     — Tiling 数据结构定义             │
├─────────────────────────────────────────────────────────────────┤
│  op_kernel/q_matmul.cpp        — Kernel 实现 (AICore 侧)       │
│  op_kernel/q_matmul_tiling.h   — Kernel 侧 Tiling 宏与结构体   │
├─────────────────────────────────────────────────────────────────┤
│  CMakePresets.json / build.sh  — 编译与打包                    │
└─────────────────────────────────────────────────────────────────┘
```

### 算子原型

| 参数 | 类型 | 格式 | 说明 |
|------|------|------|------|
| a | FP16 / FP32 | ND | 输入四元数矩阵 A, `(M, K, 4)` |
| b | FP16 / FP32 | ND | 输入四元数矩阵 B, `(K, N, 4)` |
| bExpand | FP16 / FP32 | ND | 展开后的 B 矩阵, `(4K, 4N)` - **由调用侧预先分配** |
| c | FP16 / FP32 | ND | 输出四元数矩阵 C, `(M, N, 4)` |

### 关于 `bExpand` 辅助输入

本算子需要一个额外的 `bExpand` 输入张量，这是因为 **310B 不支持 `UserWorkspace`**（片上可编程内存的运行时动态分配），因此展开后的 B 矩阵必须由调用侧在 Global Memory 中预先分配并提供。该张量占用 `4K × 4N × dtypeSize` 字节。

### Kernel 执行流程

```
Process()
  ├── ExpandBSet()
  │    遍历 B 的每个 (i, j) 位置，读取 4 个分量
  │    根据 4×4 映射表将分量写入 bExpandGlobal
  │    └── 结果: bExpandGlobal 形状 (4K, 4N)
  │
  ├── DataCacheCleanAndInvalid()
  │    刷新 bExpand 的 L1/L2 缓存，保证 Cube 读到最新数据
  │
  └── DoMatmul()
        Init Matmul 对象，配置 Cube Tiling
        SetTensorA(aGlobal) — A 视为 (M, 4K)
        SetTensorB(bExpGlobal) — B 视为 (4K, 4N)
        IterateAll(cGlobal) — 单次 Cube 矩阵乘
        └── 结果: cGlobal 形状 (M, 4N) → 视图为 (M, N, 4)
```

### Tiling

- 使用 AscendC 标准 `MatmulApiTiling` 计算 Cube 分块参数
- 数据类型支持 FP16 和 FP32，自动适配
- 单核执行（`SetBlockDim(1)`）；对于大规模矩阵，可通过修改 tiling 策略支持多核

---

## 部署与测试

### 1. 配置环境

编辑 `CMakePresets.json`，将 `ASCEND_CANN_PACKAGE_PATH` 改为实际 CANN 安装路径：

```json
"ASCEND_CANN_PACKAGE_PATH": {
    "type": "PATH",
    "value": "/path/to/Ascend/ascend-toolkit/latest"
}
```

### 2. 编译打包

```bash
cd QMatmul  # 算子工程根目录
source /path/to/Ascend/ascend-toolkit/latest/bin/setenv.bash
export DDK_PATH=/path/to/Ascend/ascend-toolkit/latest
export NPU_HOST_LIB=/path/to/Ascend/ascend-toolkit/latest/lib64

bash build.sh
```

编译完成后，安装包位于 `build_out/` 目录下。

### 3. 部署算子

```bash
cd build_out
chmod +x custom_opp_ubuntu_aarch64.run
./custom_opp_ubuntu_aarch64.run
```

### 4. 运行测试

`example/AclNNInvocationNaive/` 目录下包含一个完整的 AclNN 直调测试工程：

```bash
cd example/AclNNInvocationNaive
bash run.sh
```

流程：生成随机数据 → 编译测试程序 → 运行算子 → 对比 golden 数据验证精度。

预期输出：
```
=== Precision verification PASSED!
```

---

## 项目结构

```
QMatmul/
├── CMakeLists.txt                # 顶层 CMake 配置
├── CMakePresets.json             # CMake 预设（CANN 路径、编译选项）
├── QMatmul.json                  # 算子描述文件（注册到 OPP）
├── build.sh                      # 编译打包入口脚本
├── cmake/                        # CMake 辅助模块
│   ├── config.cmake
│   ├── func.cmake
│   ├── intf.cmake
│   ├── makeself.cmake
│   ├── device_task.cmake
│   └── util/                     # 编译工具脚本
├── op_host/                      # Host 侧（算子注册与 Tiling）
│   ├── q_matmul.cpp              #  算子原型、Shape 推导、Tiling
│   └── q_matmul_tiling.h         #  TilingData 结构定义
├── op_kernel/                    # Device 侧（AICore Kernel）
│   ├── q_matmul.cpp              #  四元数矩阵乘 Kernel 实现
│   └── q_matmul_tiling.h         #  Kernel 侧 Tiling 宏
├── example/                      # 示例工程
│   └── AclNNInvocationNaive/     #  AclNN 直调测试
│       ├── main.cpp              #  测试程序
│       ├── CMakeLists.txt
│       ├── gen_data.py           #  生成随机测试数据
│       ├── verify_result.py      #  精度验证脚本
│       ├── run.sh                #  一键运行脚本
│       ├── input/                #  输入数据目录
│       └── output/               #  输出数据与 golden 目录
└── scripts/                      # 安装/升级脚本
    ├── install.sh
    └── upgrade.sh
```

---

## 性能说明

通过将 16 次小矩阵乘法 + 12 次向量操作合并为**单次 Cube 矩阵乘法**，本算子：

- **最大化 Cube 利用率**：一次大矩阵乘比多次小矩阵乘更充分发挥 Cube 的脉动阵列
- **减少数据搬移**：A 无需重排（零拷贝视图），B 的展开一次完成
- **减少中间结果带宽**：无多次写回/读入 intermediate 结果的开销

对于 M、K、N 较大的场景，性能优势尤为显著。

---

## 开发算子信息

- **算子名称**：QMatmul
- **适配平台**：Ascend 310B
- **编程框架**：Ascend C（基于 CANN）
- **数据类型**：FP16 / FP32
- **核数**：单核（可通过扩展 tiling 支持多核）