cubecl-llvm 0.11.0-pre.4

LLVM compiler for CubeCL
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

//! NVPTX synchronization. CTA and warp barriers also order memory.

use crate::prelude::*;

const BARRIER_CTA: &str = "llvm.nvvm.barrier.cta.sync.aligned.all";

const BARRIER_WARP: &str = "llvm.nvvm.bar.warp.sync";

/// Barrier resource reserved for cube synchronization.
const BARRIER_ID: i32 = 0;

/// All warp lanes participate.
const FULL_MASK: i32 = -1;

fn barrier(scope: &Scope, name: &str, operand: i32) {
    let void_ty = VoidType::get(scope.ctx_mut()).into();
    let operand = i32_const_op(scope.ctx_mut(), operand);
    scope.register(&operand);
    let arg = operand.get_result(scope.ctx());
    let op = call_op(scope.ctx_mut(), name, void_ty, vec![arg]);
    scope.register(&op);
}

pub fn lower_sync_plane(scope: &Scope) {
    barrier(scope, BARRIER_WARP, FULL_MASK);
}

pub fn lower_sync_cube(scope: &Scope) {
    barrier(scope, BARRIER_CTA, BARRIER_ID);
}