Make lossless_cast more aggressive

halide · Jun 11, 2020 · 394536f · 394536f
1 parent 8cddb2e
commit 394536f
Showing 1 changed file with 36 additions and 0 deletions.
diff --git a/src/IROperator.cpp b/src/IROperator.cpp
@@ -429,6 +429,42 @@ Expr lossless_cast(Type t, Expr e) {
         }
     }
 
+    if ((t.is_int() || t.is_uint()) && t.bits() >= 16) {
+        if (const Add *add = e.as<Add>()) {
+            // If we can losslessly narrow the args even more
+            // aggressively, we're good.
+            // E.g. lossless_cast(uint16, (uint32)(some_u8) + 37)
+            // = (uint16)(some_u8) + 37
+            Expr a = lossless_cast(t.with_bits(t.bits() / 2), add->a);
+            Expr b = lossless_cast(t.with_bits(t.bits() / 2), add->b);
+            if (a.defined() && b.defined()) {
+                return cast(t, a) + cast(t, b);
+            } else {
+                return Expr();
+            }
+        }
+
+        if (const Sub *sub = e.as<Sub>()) {
+            Expr a = lossless_cast(t.with_bits(t.bits() / 2), sub->a);
+            Expr b = lossless_cast(t.with_bits(t.bits() / 2), sub->b);
+            if (a.defined() && b.defined()) {
+                return cast(t, a) + cast(t, b);
+            } else {
+                return Expr();
+            }
+        }
+
+        if (const Mul *mul = e.as<Mul>()) {
+            Expr a = lossless_cast(t.with_bits(t.bits() / 2), mul->a);
+            Expr b = lossless_cast(t.with_bits(t.bits() / 2), mul->b);
+            if (a.defined() && b.defined()) {
+                return cast(t, a) * cast(t, b);
+            } else {
+                return Expr();
+            }
+        }
+    }
+
     return Expr();
 }