cubecl_ir/interfaces/
uniformity.rs

1use pliron::builtin::ops::ConstantOp;
2use smallvec::SmallVec;
3
4use crate::{
5    PropagatesUniformity,
6    interfaces::control_flow::{
7        RegionBranchOpInterface, RegionBranchTerminatorOpInterface, RegionSuccessor,
8    },
9    prelude::*,
10};
11
12#[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Debug, Default)]
13pub enum Uniformity {
14    #[default]
15    Uninitialized = 4,
16    Device = 3,
17    Cube = 2,
18    Plane = 1,
19    None = 0,
20}
21
22#[op_interface]
23pub trait UniformOpInterface {
24    verify_op_succ!();
25    fn uniformity(&self, ctx: &Context, operands: &[Uniformity]) -> Uniformity;
26}
27
28PropagatesUniformity!(ConstantOp);
29
30#[op_interface]
31pub trait UniformRegionOpInterface: RegionBranchOpInterface + SingleBlockRegionInterface {
32    verify_op_succ!();
33    fn result_uniformity(&self, ctx: &Context, operands: &[Uniformity]) -> Uniformity;
34    fn entry_successor_region_uniformity(
35        &self,
36        ctx: &Context,
37        operands: &[Uniformity],
38    ) -> Vec<Uniformity>;
39}
40
41#[op_interface]
42pub trait UniformRegionTerminatorOpInterface: RegionBranchTerminatorOpInterface {
43    verify_op_succ!();
44    fn all_successor_regions(&self, ctx: &Context) -> Vec<RegionSuccessor> {
45        let operands = self
46            .get_operation()
47            .deref(ctx)
48            .operands()
49            .map(|_| None)
50            .collect::<SmallVec<[_; 8]>>();
51        self.successor_regions(ctx, &operands)
52    }
53    fn successor_region_uniformity(
54        &self,
55        ctx: &Context,
56        operands: &[Uniformity],
57    ) -> Vec<Uniformity>;
58}
cubecl_ir/interfaces/uniformity.rs

cubecl_ir/interfaces/
uniformity.rs
