[AI][昇腾950] Scalar 原子操作 (Atomic) 与注意事项

mlidongfeng

287人浏览 · 2026-06-27 07:50:45

mlidongfeng · 2026-06-27 07:50:45 发布

第一章：原子操作总览

1.1 为什么需要原子操作

问题场景: 多个 Scalar 核同时修改同一 GM 地址

Core 0: LD X0, [counter]    → X0 = 100
Core 1: LD X0, [counter]    → X0 = 100  (都读到了 100)
Core 0: ADD X0, X0, #1      → X0 = 101
Core 1: ADD X0, X0, #1      → X0 = 101  (各自加 1)
Core 0: ST X0, [counter]    → [counter] = 101
Core 1: ST X0, [counter]    → [counter] = 101  (应该是 102!)

两次 +1, 但结果只 +1 了 → 数据竞争 (Race Condition)

原子操作解决:
Core 0: ATOM.add.u32 Xm, [counter], Xt  → [counter] = 101 (原子完成)
Core 1: ATOM.add.u32 Xm, [counter], Xt  → [counter] = 102 (原子完成)

1.2 四类原子指令定位

┌─────────────────────────────────────────────────────────┐
│              Scalar 核原子操作分类                         │
├──────────────┬──────────┬───────────┬───────────────────┤
│ 指令         │ 总线     │ 地址空间   │ 返回旧值?         │
├──────────────┼──────────┼───────────┼───────────────────┤
│ ATOM         │ 系统总线  │ GM (OUT)  │  返回旧值        │
│ RED          │ 系统总线  │ GM (OUT)  │ 不返回            │
└──────────────┴──────────┴───────────┴───────────────────┘

第二章：ATOM — 全局内存原子操作

2.1 指令格式

T AtomicCAS(__gm__ T *addr, T compare, T val);
T AtomicExch(__gm__ T *addr, T val);
T AtomicSub(__gm__ T *addr, T val);
T AtomicAdd(__gm__ T *addr, T val);
T AtomicMin(__gm__ T *addr, T val);
T AtomicMax(__gm__ T *addr, T val);

参数说明:
  T — 数据类型: u32/s32/u64/s64

2.2 各操作语义详解

.cas — Compare-And-Swap (核心同步原语)

// C 语言等价语义
uint32_t AtomicCAS(uint32_t *addr, uint32_t expected, uint32_t desired) {
    uint32_t old = *addr;
    if (old == expected) {
        *addr = desired;    // 匹配则写入 desired
    }
    return old;             // 总是返回旧值
}

.exch — 原子交换

// C 语言等价语义
uint32_t AtomicExch(uint32_t *addr, uint32_t new_val) {
    uint32_t old = *addr;
    *addr = new_val;
    return old;
}

.add — 原子加

// C 语言等价语义
uint32_t AtomicAdd(uint32_t *addr, uint32_t val) {
    uint32_t old = *addr;
    *addr = old + val;
    return old;
}

.min / .max — 原子最小/最大值

// C 语言等价语义
uint32_t AtomicMax(uint32_t *addr, uint32_t val) {
    uint32_t old = *addr;
    *addr = (val > old) ? val : old;
    return old;
}

第三章：RED — 全局内存归约操作

3.1 指令格式

void AtomicAdd(__gm__ T *addr, T val);
void AtomicMin(__gm__ T *addr, T val);
void AtomicMax(__gm__ T *addr, T val);


参数说明:
  .T — 数据类型: u32/s32/f16/bf16/f32

关键区别:
  - 不返回旧值
  - 支持浮点类型 (f16/bf16/f32) ← ATOM 不支持!

3.3 语义

// *addr = *addr op Xm  (直接在目标地址上归约, 不返回旧值)
*Xn = *Xn op Xm;

// 例: RED.add.f32 [Xn], Xm
// *Xn = *Xn + Xm  (原子浮点累加, 不返回旧值)

3.4 ATOM vs RED 选择指南

┌─────────────────────────────────────────────────┐
│            ATOM vs RED 选择决策                   │
├─────────────────────────────────────────────────┤
│                                                 │
│  需要读取旧值?                                    │
│    ├── YES → ATOM                               │
│    │    适用: 锁实现, 需要判断是否成功的场景       │
│    │                                             │
│    └── NO → 数据类型是浮点?                      │
│         ├── YES → RED (唯一能做浮点原子归约的)    │
│         │    适用: AllReduce float, 全局浮点累加  │
│         │                                        │
│         └── NO (整数)                            │
│              ├── 需要知道旧值 → ATOM             │
│              └── 纯归约 → RED (省掉 Xt 寄存器)   │
└─────────────────────────────────────────────────┘

第四章：实战编程模式

4.1 Spin Lock (自旋锁)

// ══════════════════════════════════════════════════
// 自旋锁: 用 CAS 实现
// lock_addr 指向一个 u32 变量: 0=解锁, 1=加锁
// ══════════════════════════════════════════════════

// 获取锁
  // expected = 0 (未锁),  desired  = 1 (加锁)
  while( true){  
      if( AtomicCAS(gm, 0, 1) == 0  ) { //返回0， 表示加锁成功
         break;
      }
  }
//获取到锁，处理业务。

//释放锁
  AtomicSub(gm, 1);   // 写 0 释放锁
  DSB(ALL)             // 确保后续指令可见,  按需写，不是必须的
  // 清 dcache,  这个是不需要的， 因为atomic 是 bypass DCache

4.2 多核 Barrier (栅栏同步)

// ══════════════════════════════════════════════════
// Barrier: N 个 Core 都到达后才能继续
// barrier_counter: 初始为 0, 每个到达的 Core 原子 +1
// ══════════════════════════════════════════════════

// 每个 Core 到达 barrier 时:
  auto barrier_counter = AtomicAdd(gm, 1);  //  barrier_counter 旧值 (= 之前已到达的 Core 数)
  if( barrier_counter == N ) {
    sfv(); //唤醒
  }
  else {
    while( barrier_counter != N ){
      wfe();  //休眠    
      barrier_counter = AtomicAdd(gm, 0);
    }
  }

第五章：Cache 一致性维护 (核心关键)

5.1 原子操作的 Cache 行为

┌─────────────────────────────────────────────────────────┐
│           ATOM/RED 的 Cache 行为                         │
├─────────────────────────────────────────────────────────┤
│                                                         │
│  ATOM/RED 旁路 DCache  :                                │
│                                                         │
│    CPU Core                                             │
│      │                                                  │
│      ├── LD / ST ──→ DCache ──→ L2 ──→ GM               │
│      │                 ↑ 可能缓存旧值                    │
│      │                                                  │
│      └── ATOM/RED ────────────L2 ──────→ GM (直接!)     │
│                   旁路DCache, L2可配                    │
│                                                         │
│  问题:                                                  │
│    T0: LD X0, [addr]    → X0 = 100 (cache 中缓存 100)  │
│    T1: ATOM.add [addr]  → GM 中 addr 变成 101           │
│    T2: LD X0, [addr]    → X0 = 100 !! (cache hit 旧值!)│
│                                                         │
│  原子操作不会自动清理 dcache!                             │
└─────────────────────────────────────────────────────────┘

5.2 一致性维护协议

规则 1: ATOM/RED 之后, 如果需要通过 LD 读取同一地址, 必须清 cache
 
规则 2: 通过 LD 读取后, 如果之后要 ATOM 同一地址, 也建议清 cache
 
规则 3: 多核场景, 写端 ATOM 后通知读端, 读端必须先清 cache

5.3 饱和模式与 FP16 原子

当 CTRL[8:6] = 3'b010 (f16 原子) 时:

CTRL[48] = 0 (饱和模式):
  FP16 原子操作中:
    溢出 → 钳位到 ±65504
    NaN  → 钳位为 0
  适合推理场景

CTRL[48] = 1 (IEEE 模式):
  FP16 原子操作中:
    溢出 → ±INF
    NaN  → 正常传播
  适合训练场景

⚠ 警告: 修改 CTRL[48] 前, 如果之前执行过 FP16 的隐式原子存储,
  必须先执行 DCCI Xn, 1, #ATOMIC 清理残留 cache!

第六章：注意事项总结

6.1 硬约束 (违反会异常或数据错误)

编号	约束	后果	正确做法
A1	ATOM/RED 禁止访问栈地址	异常	只对 GM (OUT) 地址使用
A2	地址 bit[63:49] ≠ 0	地址溢出异常	确保地址在合法范围内
A3	地址未对齐到 type 大小	对齐异常	u32→4B 对齐, u64→8B 对齐
A4	cacheable 与 non-cacheable 地址相隔 < 4KB	数据不一致	地址规划预留 4KB 间隔
A5	ATOM/RED 计入 DSB.ALL 和 DSB.DDR	需等待完成	DSB 后才能确认操作完成

6.2 一致性约束 (违反会数据错误)

编号	约束	后果	正确做法
C1	ATOM/RED 后不清理 cache	LD 读到旧值	ATOM → DSB → DCCI ATOMIC
C2	多核 ATOM 后不通知	对端不知道有新数据	ATOM → DSB → DCCI → SET_CROSS_CORE
C3	读取端不清理 cache	读到本地缓存的旧值	WAIT_FLAG → DSB → DCCI → LD
C4	切换 CTRL[48] 前不清理 ATOMIC cache	FP16 饱和/IEEE 模式切换不一致	先 DCCI ATOMIC 再改 CTRL[48]
C5	SSBUF 上使用原子操作	SSBUF 不支持原子	只对 GM/HSCB 使用原子操作

6.3 性能注意事项

编号	注意事项	影响	建议
P1	ATOM 走系统总线, 延迟高 (~200ns)	频繁 ATOM 会成为瓶颈	减少原子操作频率, 批量累加后一次 ATOM
P2	ATOM.cas 自旋等待消耗总线带宽	锁竞争激烈时性能下降	临界区尽量短; 考虑公平锁
P3	每次都 DSB + DCCI 开销大	增加 10-20 周期	批量原子操作后统一 DSB + DCCI
P5	RED 比 ATOM 少返回旧值	少一个寄存器操作	不需要旧值时优先用 RED
P4	ATOM.u64 比 ATOM.u32 慢	64bit 原子操作带宽翻倍	能用 u32 就不用 u64

6.4 操作速查矩阵

需求	指令	数据类型	是否返回旧值	总线
原子加 (整数)	ATOM.add	u32/s32/u64/s64	Yes	系统总线
原子加 (浮点)	RED.add	f16/bf16/f32	No	系统总线
原子最大 (整数)	ATOM.max	u32/s32/u64/s64	Yes	系统总线
原子最大 (浮点)	RED.max	f16/bf16/f32	No	系统总线
原子最小 (整数)	ATOM.min	u32/s32/u64/s64	Yes	系统总线
原子最小 (浮点)	RED.min	f16/bf16/f32	No	系统总线
互斥锁	ATOM.cas	u32/u64	Yes	系统总线
原子交换	ATOM.exch	u32/u64	Yes	系统总线