cubecl_core/post_processing/
checked_io.rs

1use alloc::vec;
2
3use alloc::string::String;
4use cubecl_ir::{
5    NamedRewrite, Scope, dialect::memory::IndexOp, prelude::*, settings::ExecutionMode,
6    types::RuntimeArrayType,
7};
8
9use crate::io::*;
10
11pub type CheckedIoPass = MatchRewritePass<CheckedIo>;
12
13#[derive(new, NamedRewrite)]
14pub struct CheckedIo {
15    mode: ExecutionMode,
16    kernel_name: String,
17}
18
19impl MatchRewrite for CheckedIo {
20    fn r#match(&mut self, ctx: &Context, op: Ptr<Operation>) -> bool {
21        Operation::get_op::<IndexOp>(op, ctx)
22            .is_some_and(|it| it.checked(ctx) && is_runtime_array(ctx, it.base(ctx)))
23    }
24
25    fn rewrite(
26        &mut self,
27        ctx: &mut Context,
28        rewriter: &mut DialectConversionRewriter,
29        op: Ptr<Operation>,
30    ) -> Result<()> {
31        let index = Operation::get_op::<IndexOp>(op, ctx).unwrap();
32
33        let scope = Scope::from_context_and_inserter(ctx, rewriter);
34
35        let new_value = match self.mode {
36            ExecutionMode::Checked => {
37                expand_checked_index(&scope, index.base(ctx), index.index(ctx))
38            }
39            ExecutionMode::Validate => {
40                expand_validate_index(&scope, index.base(ctx), index.index(ctx), &self.kernel_name)
41            }
42            ExecutionMode::Unchecked => index.get_result(ctx),
43        };
44        rewriter.replace_operation_with_values(ctx, index.get_operation(), vec![new_value]);
45        Ok(())
46    }
47}
48
49fn is_runtime_array(ctx: &Context, value: Value) -> bool {
50    let ty = value.get_type(ctx).deref(ctx);
51    ty.downcast_ref::<RuntimeArrayType>().is_some()
52}
cubecl_core/post_processing/checked_io.rs

cubecl_core/post_processing/
checked_io.rs
