Small opt

jafioti · Apr 26, 2024 · 45dc0f6 · 45dc0f6
1 parent 358dceb
commit 45dc0f6
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 22 deletions.
diff --git a/crates/luminal_metal/src/elementwise_fusion.rs b/crates/luminal_metal/src/elementwise_fusion.rs
@@ -8,12 +8,9 @@ use metal_rs::{
     ComputePipelineState, Device, MTLResourceOptions,
 };
 
-use luminal::{
-    op::{InputTensor, Operator},
-    prelude::{
-        petgraph::{visit::EdgeRef, Direction},
-        *,
-    },
+use luminal::prelude::{
+    petgraph::{visit::EdgeRef, Direction},
+    *,
 };
 
 use crate::{
@@ -384,11 +381,15 @@ impl<T: MetalFloat> Compiler for ElementwiseFusionCompiler<T> {
                     *subexp = re
                         .replace_all(
                             subexp,
-                            &format!(
-                                "({} != 0 ? (float)input{i}[{}] : 0.0)$1",
-                                expr_to_metal_string(val_exp.clone()),
-                                expr_to_metal_string(ind_exp.clone())
-                            ),
+                            &if *val_exp != true {
+                                format!(
+                                    "({} != 0 ? (float)input{i}[{}] : 0.0)$1",
+                                    expr_to_metal_string(val_exp),
+                                    expr_to_metal_string(ind_exp)
+                                )
+                            } else {
+                                format!("(float)input{i}[{}]$1", expr_to_metal_string(ind_exp))
+                            },
                         )
                         .to_string();
                 }
@@ -410,7 +411,7 @@ impl<T: MetalFloat> Compiler for ElementwiseFusionCompiler<T> {
                     if val_exp != true {
                         *subexp = format!(
                             "(({} != 0) ? {subexp} : 0.0)",
-                            expr_to_metal_string(val_exp)
+                            expr_to_metal_string(&val_exp)
                         );
                     }
                 }

diff --git a/crates/luminal_metal/src/lib.rs b/crates/luminal_metal/src/lib.rs
@@ -395,9 +395,9 @@ fn render_dyn_dim_inputs(shapes: &[ShapeTracker], offset: usize) -> (Vec<char>,
     )
 }
 
-fn expr_to_metal_string(expr: BigExpression) -> String {
+fn expr_to_metal_string(expr: &BigExpression) -> String {
     let mut symbols = vec![];
-    for term in expr.terms {
+    for term in expr.terms.clone() {
         let new_symbol = match term {
             Term::Num(n) => n.to_string(),
             Term::Var(c) => {
@@ -440,8 +440,8 @@ fn expr_to_metal_string(expr: BigExpression) -> String {
 
 fn get_idx_valid_exps(shape: ShapeTracker) -> (String, String) {
     (
-        expr_to_metal_string(shape.index_expression()),
-        expr_to_metal_string(shape.valid_expression()),
+        expr_to_metal_string(&shape.index_expression()),
+        expr_to_metal_string(&shape.valid_expression()),
     )
 }
 

diff --git a/crates/luminal_symbolic/src/expression.rs b/crates/luminal_symbolic/src/expression.rs
@@ -146,9 +146,12 @@ impl<S: ExpressionStorage> GenericExpression<S> {
     /// Maximum
     pub fn max<E: Into<Self>>(self, rhs: E) -> Self {
         let mut rhs = rhs.into();
-        if rhs == self {
+        if rhs == self || rhs == 0 {
             return self;
         }
+        if self == 0 {
+            return rhs;
+        }
         rhs.terms.extend(self.terms);
         rhs.terms.push(Term::Max);
         rhs.simplify()

diff --git a/src/shape/tracker.rs b/src/shape/tracker.rs
@@ -108,11 +108,6 @@ impl ShapeTracker {
         let mut ind_expr = BigExpression::from(0); // The final index expression
         let mut current_elem_size = BigExpression::from(1); // Keep track of the size of each element of the current dim (last dim elem size: 1)
 
-        // For combined dims
-        // divide by last dims (smallest) element size
-        // mod by combined dim size
-        // multiply by last dims (smallest) stride
-
         // Loop through all dims in reverse order
         for i in shape.indexes.into_iter().rev() {
             // Get logical dimension size with padding and mask