cubecl_ir/
pipeline.rs

1use crate::{Instruction, TypeHash};
2use core::fmt::Display;
3
4use crate::OperationReflect;
5
6use super::Variable;
7
8#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
9#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationReflect)]
10#[operation(opcode_name = PipelineOpCode)]
11/// Operations available on a pipeline
12pub enum PipelineOps {
13    /// Copy source to destination
14    MemCopyAsync {
15        pipeline: Variable,
16        source: Variable,
17        destination: Variable,
18    },
19    /// Reserves a specific stage for the producer to work on.
20    ProducerAcquire { pipeline: Variable },
21    /// Signals that the producer is done and the stage is ready for the consumer.
22    ProducerCommit { pipeline: Variable },
23    /// Waits until the producer has finished with the stage.
24    ConsumerWait { pipeline: Variable },
25    /// Frees the stage after the consumer is done using it.
26    ConsumerRelease { pipeline: Variable },
27}
28
29impl Display for PipelineOps {
30    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
31        match self {
32            PipelineOps::MemCopyAsync {
33                pipeline,
34                source,
35                destination,
36            } => write!(
37                f,
38                "mem_copy_async({pipeline}, source: {source}, destination: {destination})",
39            ),
40            PipelineOps::ProducerAcquire { pipeline } => write!(f, "producer_acquire({pipeline})"),
41            PipelineOps::ProducerCommit { pipeline } => write!(f, "producer_commit({pipeline})"),
42            PipelineOps::ConsumerWait { pipeline } => write!(f, "consumer_wait({pipeline})"),
43            PipelineOps::ConsumerRelease { pipeline } => write!(f, "consumer_release({pipeline})"),
44        }
45    }
46}
47
48impl From<PipelineOps> for Instruction {
49    fn from(value: PipelineOps) -> Self {
50        Instruction::no_out(value)
51    }
52}
cubecl_ir/pipeline.rs

cubecl_ir/
pipeline.rs
