Implement inclusive ranges

tracel-ai · Sep 9, 2024 · 6a94008 · 6a94008
1 parent 02bc447
commit 6a94008
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 2 deletions.
diff --git a/crates/cubecl-core/src/ir/processing.rs b/crates/cubecl-core/src/ir/processing.rs
@@ -246,6 +246,9 @@ impl ScopeProcessing {
                 Branch::RangeLoop(op) => {
                     sanitize_constant_scalar_ref_elem(&mut op.start, Elem::UInt);
                     sanitize_constant_scalar_ref_elem(&mut op.end, Elem::UInt);
+                    if let Some(step) = &mut op.step {
+                        sanitize_constant_scalar_ref_elem(step, Elem::UInt);
+                    }
                 }
                 _ => {
                     // Nothing to do.

diff --git a/crates/cubecl-cuda/src/compiler/base.rs b/crates/cubecl-cuda/src/compiler/base.rs
@@ -295,6 +295,7 @@ impl CudaCompiler {
                 start: self.compile_variable(range_loop.start),
                 end: self.compile_variable(range_loop.end),
                 step: range_loop.step.map(|it| self.compile_variable(it)),
+                inclusive: range_loop.inclusive,
                 instructions: self.compile_scope(&mut range_loop.scope),
             }),
             gpu::Branch::Loop(mut op) => instructions.push(Instruction::Loop {

diff --git a/crates/cubecl-cuda/src/compiler/instruction.rs b/crates/cubecl-cuda/src/compiler/instruction.rs
@@ -50,6 +50,7 @@ pub enum Instruction {
         start: Variable,
         end: Variable,
         step: Option<Variable>,
+        inclusive: bool,
         instructions: Vec<Self>,
     },
     Loop {
@@ -188,15 +189,17 @@ impl Display for Instruction {
                 start,
                 end,
                 step,
+                inclusive,
                 instructions,
             } => {
                 let increment = step
                     .map(|step| format!("{i} += {step}"))
                     .unwrap_or_else(|| format!("++{i}"));
+                let cmp = if *inclusive { "<=" } else { "<" };
 
                 f.write_fmt(format_args!(
                     "
-for (uint {i} = {start}; {i} < {end}; {increment}) {{
+for (uint {i} = {start}; {i} {cmp} {end}; {increment}) {{
 "
                 ))?;
                 for instruction in instructions {

diff --git a/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs b/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs
@@ -369,6 +369,7 @@ impl WgslCompiler {
                     start: self.compile_variable(range_loop.start),
                     end: self.compile_variable(range_loop.end),
                     step: range_loop.step.map(|it| self.compile_variable(it)),
+                    inclusive: range_loop.inclusive,
                     instructions: self.compile_scope(&mut range_loop.scope),
                 })
             }

diff --git a/crates/cubecl-wgpu/src/compiler/wgsl/instructions.rs b/crates/cubecl-wgpu/src/compiler/wgsl/instructions.rs
@@ -182,6 +182,7 @@ pub enum Instruction {
         start: Variable,
         end: Variable,
         step: Option<Variable>,
+        inclusive: bool,
         instructions: Vec<Instruction>,
     },
     And {
@@ -531,16 +532,18 @@ impl Display for Instruction {
                 start,
                 end,
                 step,
+                inclusive,
                 instructions,
             } => {
                 let increment = step
                     .as_ref()
                     .map(|step| format!("{i} += {step}"))
                     .unwrap_or_else(|| format!("{i}++"));
+                let cmp = if *inclusive { "<=" } else { "<" };
 
                 f.write_fmt(format_args!(
                     "
-for (var {i}: u32 = {start}; {i} < {end}; {increment}) {{
+for (var {i}: u32 = {start}; {i} {cmp} {end}; {increment}) {{
 "
                 ))?;
                 for instruction in instructions {