cubecl-ir 0.11.1

Intermediate representation for CubeCL
Documentation 
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
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58

use pliron::builtin::ops::ConstantOp;
use smallvec::SmallVec;

use crate::{
    PropagatesUniformity,
    interfaces::control_flow::{
        RegionBranchOpInterface, RegionBranchTerminatorOpInterface, RegionSuccessor,
    },
    prelude::*,
};

#[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Debug, Default)]
pub enum Uniformity {
    #[default]
    Uninitialized = 4,
    Device = 3,
    Cube = 2,
    Plane = 1,
    None = 0,
}

#[op_interface]
pub trait UniformOpInterface {
    verify_op_succ!();
    fn uniformity(&self, ctx: &Context, operands: &[Uniformity]) -> Uniformity;
}

PropagatesUniformity!(ConstantOp);

#[op_interface]
pub trait UniformRegionOpInterface: RegionBranchOpInterface + SingleBlockRegionInterface {
    verify_op_succ!();
    fn result_uniformity(&self, ctx: &Context, operands: &[Uniformity]) -> Uniformity;
    fn entry_successor_region_uniformity(
        &self,
        ctx: &Context,
        operands: &[Uniformity],
    ) -> Vec<Uniformity>;
}

#[op_interface]
pub trait UniformRegionTerminatorOpInterface: RegionBranchTerminatorOpInterface {
    verify_op_succ!();
    fn all_successor_regions(&self, ctx: &Context) -> Vec<RegionSuccessor> {
        let operands = self
            .get_operation()
            .deref(ctx)
            .operands()
            .map(|_| None)
            .collect::<SmallVec<[_; 8]>>();
        self.successor_regions(ctx, &operands)
    }
    fn successor_region_uniformity(
        &self,
        ctx: &Context,
        operands: &[Uniformity],
    ) -> Vec<Uniformity>;
}