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
31
32
33
34
35

//! AMDGPU synchronization.

use crate::prelude::*;

const S_BARRIER: &str = "llvm.amdgcn.s.barrier";

const WAVE_BARRIER: &str = "llvm.amdgcn.wave.barrier";

fn fence(scope: &Scope, sync_scope: &str, ordering: AtomicOrderingAttr) {
    let fence = llvm::FenceOp::new(
        scope.ctx_mut(),
        ordering,
        SyncScopeAttr::NamedScope(sync_scope.into()),
    );
    scope.register(&fence);
}

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

/// Wavefront synchronization requires memory ordering and a scheduling barrier.
pub fn lower_sync_plane(scope: &Scope) {
    fence(scope, "wavefront", AtomicOrderingAttr::AcqRel);
    barrier(scope, WAVE_BARRIER);
}

/// Cube synchronization requires memory ordering and a workgroup barrier.
pub fn lower_sync_cube(scope: &Scope) {
    fence(scope, "workgroup", AtomicOrderingAttr::Release);
    barrier(scope, S_BARRIER);
    fence(scope, "workgroup", AtomicOrderingAttr::Acquire);
}