# 昇腾算子之四元数矩阵乘法

**Repository Path**: cluo29/qmatmul

## Basic Information

- **Project Name**: 昇腾算子之四元数矩阵乘法
- **Description**: 昇腾算子之四元数矩阵乘法
- **Primary Language**: Unknown
- **License**: MulanPSL-2.0
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-08-19
- **Last Updated**: 2026-08-19

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# 昇腾算子之四元数矩阵乘法

## 开发者
https://gitee.com/Nicet

大湾区大学信息科学技术学院，胡人

大湾区大学信息科学技术学院，罗除，https://gitee.com/cluo29

## 课题承担方
大湾区大学信息科学技术学院，罗除，https://gitee.com/cluo29

## 课题委托方
人工智能与数字经济广东省实验室（深圳）

## 项目支撑
由人工智能与数字经济广东省实验室（深圳） 开放课题资助（编号：GML-KF-24-01）
This research was financially supported by the Open Research Fund from Guangdong Laboratory of Artificial Intelligence and Digital Economy (SZ), under Grant No.GML-KF-24-01.

# QMatmul 算子

## 算子介绍

本算子实现四元数张量矩阵乘法。

A shape: (M, K, 4)
B shape: (K, N, 4)
计算得到 C shape: (M, N, 4)

算子内部实现是将矩阵 B 局部展开，下面是验证局部展开策略的可行性：

```py
import torch

# ==========================================
# 1. 原始实现 (用于对标结果)
# ==========================================
def quaternion_matmul_tensor(A, B):
    A_r, A_i, A_j, A_k = torch.unbind(A, dim=-1)
    B_r, B_i, B_j, B_k = torch.unbind(B, dim=-1)

    C_r = A_r @ B_r - A_i @ B_i - A_j @ B_j - A_k @ B_k
    C_i = A_r @ B_i + A_i @ B_r + A_j @ B_k - A_k @ B_j
    C_j = A_r @ B_j - A_i @ B_k + A_j @ B_r + A_k @ B_i
    C_k = A_r @ B_k + A_i @ B_j - A_j @ B_i + A_k @ B_r

    return torch.stack([C_r, C_i, C_j, C_k], dim=-1)

# ==========================================
# 2. Ascend C 降维展开思想的 Python 实现
# ==========================================
def quaternion_matmul_expand(A, B):
    """
    通过二维矩阵展开 (A_flat @ B_expand) 来实现四元数矩阵乘法。
    这完全模拟了 Ascend 310B NPU 上使用 Cube 单元单次计算的过程。
    """
    M, K, _ = A.shape
    _, N, _ = B.shape
    
    # 1. 内存视角转换：视 A 为 [M, 4*K] 的二维矩阵 (零拷贝)
    # 内存中的数据按 r, i, j, k 交织排列
    A_flat = A.view(M, 4 * K)
    
    # 2. 提取 B 的分量 (在 NPU 中对应 DataCopy + Gather)
    B_r, B_i, B_j, B_k = torch.unbind(B, dim=-1)
    
    # 3. 在 UB 中构造 B_expand，形状为 [4*K, 4*N]
    B_expand = torch.zeros(4 * K, 4 * N, dtype=A.dtype, device=A.device)
    
    # 将哈密顿规则映射到 4x4 的分块矩阵上
    # 这里的 0::4 语法完美模拟了 NPU Vector 单元的步长写入 (Interleave)
    
    # 第 0 组交织 (原样写入)
    B_expand[0::4, 0::4] =  B_r
    B_expand[0::4, 1::4] =  B_i
    B_expand[0::4, 2::4] =  B_j
    B_expand[0::4, 3::4] =  B_k
    
    # 第 1 组交织 (对应 i 行)
    B_expand[1::4, 0::4] = -B_i
    B_expand[1::4, 1::4] =  B_r
    B_expand[1::4, 2::4] = -B_k
    B_expand[1::4, 3::4] =  B_j
    
    # 第 2 组交织 (对应 j 行)
    B_expand[2::4, 0::4] = -B_j
    B_expand[2::4, 1::4] =  B_k
    B_expand[2::4, 2::4] =  B_r
    B_expand[2::4, 3::4] = -B_i
    
    # 第 3 组交织 (对应 k 行)
    B_expand[3::4, 0::4] = -B_k
    B_expand[3::4, 1::4] = -B_j
    B_expand[3::4, 2::4] =  B_i
    B_expand[3::4, 3::4] =  B_r


    # 4. NPU Cube 单元全速启动，执行一次标准的实数矩阵乘法
    C_flat = A_flat @ B_expand
    
    # 5. 内存视角恢复：直接把 [M, 4*N] 当作 [M, N, 4] 读出 (零拷贝)
    C_expand_result = C_flat.view(M, N, 4)
    
    return C_expand_result

# ==========================================
# 3. 正确性验证
# ==========================================
if __name__ == "__main__":
    M, K, N = 52, 43, 61
    torch.manual_seed(42)
    A = torch.randint(-2, 3, (M, K, 4), dtype=torch.float32)
    B = torch.randint(-2, 3, (K, N, 4), dtype=torch.float32)
    B = torch.arange(K*N*4, dtype=torch.float32).view(K, N, 4)

    # 运行两种方法
    C_original = quaternion_matmul_tensor(A, B)
    C_expand = quaternion_matmul_expand(A, B)

    # 验证是否一致
    is_identical = torch.allclose(C_original, C_expand, atol=1e-6)
    
    print("=================== 算法验证结果 ===================")
    print(f"原始方法结果 Shape: {C_original.shape}")
    print(f"展开方法结果 Shape: {C_expand.shape}")
    print(f"两种方法结果是否 100% 一致?  >>> {'✅ 是 (True)' if is_identical else '❌ 否 (False)'}")
    
    if is_identical:
        print("\n验证成功！通过将 B 展开为 4K x 4N 的矩阵，")
        print("我们成功将 16 次小矩阵乘 + 12 次加减法 融合成了 1 次标准矩阵乘法。")
        print("这个逻辑可以直接无缝移植到 Ascend C 的 UB 到 Cube 数据流中！")
```

## 算子版本

目前实现了两个版本的算子，算子名称均为QMatmul，分别适用于ascend310b、ascend910b，均调用一个核进行计算。

二者在算子原型上存在区别，310b版本相较于910b多了一个辅助向量，BExpand，其shape是展开后的B的shape [4*K, 4*N]，这是由于在310b上无法利用userworkspace，故需要在调用阶段提供辅助向量。

## 算子部署测试

1. 更改CMakePresets.json的ASCEND_CANN_PACKAGE_PATH到实际路径
2. 激活环境
3. 到算子目录下运行`bash build.sh`
4. 部署`build_out`目录下的run包
5. example目录下包含aclnn直调测试工程，`bash run.sh`即可查看测试结果