From 394536f5d7ffea57194294c8d5efbd805902ce88 Mon Sep 17 00:00:00 2001
From: Andrew Adams <andrew.b.adams@gmail.com>
Date: Thu, 11 Jun 2020 15:35:14 -0700
Subject: [PATCH] Make lossless_cast more aggressive

---
 src/IROperator.cpp | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/src/IROperator.cpp b/src/IROperator.cpp
index b1a8d7b78a72..f32e00b86788 100644
--- 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();
 }