Make random faster by putting the innermost var last (#6504)

abadams · web-flow · commit 00211656fd20 · 2022-01-04T08:40:23.000-08:00
* Make random 2x faster by putting the innermost var last

* Improve period of low bits of random noise

* Add new rewrite rules for quadratics

By pulling constant additions outside of quadratics, we can shave off a
few add instructions in the inner loop for random number generation,
which uses a quadratic modulo 2^32

I also removed the !overflows predicates, because rules already fail to
match if a fold overflows.

New rules formally verified.

* Make expensive_zero actually always zero
diff --git a/src/Random.cpp b/src/Random.cpp
@@ -82,6 +82,10 @@ Expr random_int(const vector<Expr> &e) {
                                rng32(Variable::make(UInt(32), name)));
         }
     }
+    // The low bytes of this have a poor period, so mix in the high bytes for
+    // two additional instructions.
+    result = result ^ (result >> 16);
+
     return result;
 }
 
@@ -101,7 +105,9 @@ class LowerRandom : public IRMutator {
     Expr visit(const Call *op) override {
         if (op->is_intrinsic(Call::random)) {
             vector<Expr> args = op->args;
-            args.insert(args.end(), extra_args.begin(), extra_args.end());
+            // Insert the free vars in reverse, so innermost vars typically end
+            // up last.
+            args.insert(args.end(), extra_args.rbegin(), extra_args.rend());
             if (op->type == Float(32)) {
                 return random_float(args);
             } else if (op->type == Int(32)) {
@@ -121,14 +127,14 @@ class LowerRandom : public IRMutator {
 
 public:
     LowerRandom(const vector<VarOrRVar> &free_vars, int tag) {
-        extra_args.emplace_back(tag);
         for (const VarOrRVar &v : free_vars) {
             if (v.is_rvar) {
                 extra_args.push_back(v.rvar);
             } else {
                 extra_args.push_back(v.var);
             }
         }
+        extra_args.emplace_back(tag);
     }
 };
 
diff --git a/src/Simplify_Mul.cpp b/src/Simplify_Mul.cpp
@@ -69,12 +69,16 @@ Expr Simplify::visit(const Mul *op, ExprInfo *bounds) {
         }
 
         if (rewrite(c0 * c1, fold(c0 * c1)) ||
-            rewrite((x + c0) * c1, x * c1 + fold(c0 * c1), !overflows(c0 * c1)) ||
-            rewrite((c0 - x) * c1, x * fold(-c1) + fold(c0 * c1), !overflows(c0 * c1)) ||
+            rewrite((x + c0) * (x + c1), x * (x + fold(c0 + c1)) + fold(c0 * c1)) ||
+            rewrite((x * c0 + c1) * (x + c2), x * (x * c0 + fold(c1 + c0 * c2)) + fold(c1 * c2)) ||
+            rewrite((x + c2) * (x * c0 + c1), x * (x * c0 + fold(c1 + c0 * c2)) + fold(c1 * c2)) ||
+            rewrite((x * c0 + c1) * (x * c2 + c3), x * (x * fold(c0 * c2) + fold(c0 * c3 + c1 * c2)) + fold(c1 * c3)) ||
+            rewrite((x + c0) * c1, x * c1 + fold(c0 * c1)) ||
+            rewrite((c0 - x) * c1, x * fold(-c1) + fold(c0 * c1)) ||
             rewrite((0 - x) * y, 0 - x * y) ||
             rewrite(x * (0 - y), 0 - x * y) ||
             rewrite((x - y) * c0, (y - x) * fold(-c0), c0 < 0 && -c0 > 0) ||
-            rewrite((x * c0) * c1, x * fold(c0 * c1), !overflows(c0 * c1)) ||
+            rewrite((x * c0) * c1, x * fold(c0 * c1)) ||
             rewrite((x * c0) * y, (x * y) * c0, !is_const(y)) ||
             rewrite(x * (y * c0), (x * y) * c0) ||
             rewrite(max(x, y) * min(x, y), x * y) ||
diff --git a/test/correctness/async_device_copy.cpp b/test/correctness/async_device_copy.cpp
@@ -9,9 +9,9 @@ Expr expensive_zero(Expr x, Expr y, Expr t, int n) {
     RDom r(0, n);
     Func a, b, c;
     Var z;
-    a(x, y, t, z) = random_int() % 1024;
-    b(x, y, t, z) = random_int() % 1024;
-    c(x, y, t, z) = random_int() % 1024;
+    a(x, y, t, z) = random_int() % 1024 + 5;
+    b(x, y, t, z) = random_int() % 1024 + 5;
+    c(x, y, t, z) = random_int() % 1024 + 5;
     return sum(select(pow(a(x, y, t, r), 3) + pow(b(x, y, t, r), 3) == pow(c(x, y, t, r), 3), 1, 0));
 }
 
diff --git a/test/correctness/simplify.cpp b/test/correctness/simplify.cpp
@@ -130,7 +130,7 @@ void check_casts() {
 
     check(cast(Float(64), 0.5f), Expr(0.5));
     check((x - cast(Float(64), 0.5f)) * (x - cast(Float(64), 0.5f)),
-          (x + Expr(-0.5)) * (x + Expr(-0.5)));
+          (cast(Float(64), x) + Expr(-1.0)) * cast(Float(64), x) + Expr(0.25));
 
     check(cast(Int(64, 3), ramp(5.5f, 2.0f, 3)),
           cast(Int(64, 3), ramp(5.5f, 2.0f, 3)));

Original file line number	Diff line number	Diff line change
`@@ -82,6 +82,10 @@ Expr random_int(const vector<Expr> &e) {`
`82`	`82`	`rng32(Variable::make(UInt(32), name)));`
`83`	`83`	`}`
`84`	`84`	`}`
	`85`	`+ // The low bytes of this have a poor period, so mix in the high bytes for`
	`86`	`+ // two additional instructions.`
	`87`	`+ result = result ^ (result >> 16);`
	`88`	`+`
`85`	`89`	`return result;`
`86`	`90`	`}`
`87`	`91`
`@@ -101,7 +105,9 @@ class LowerRandom : public IRMutator {`
`101`	`105`	`Expr visit(const Call *op) override {`
`102`	`106`	`if (op->is_intrinsic(Call::random)) {`
`103`	`107`	`vector<Expr> args = op->args;`
`104`		`- args.insert(args.end(), extra_args.begin(), extra_args.end());`
	`108`	`+ // Insert the free vars in reverse, so innermost vars typically end`
	`109`	`+ // up last.`
	`110`	`+ args.insert(args.end(), extra_args.rbegin(), extra_args.rend());`
`105`	`111`	`if (op->type == Float(32)) {`
`106`	`112`	`return random_float(args);`
`107`	`113`	`} else if (op->type == Int(32)) {`
`@@ -121,14 +127,14 @@ class LowerRandom : public IRMutator {`
`121`	`127`
`122`	`128`	`public:`
`123`	`129`	`LowerRandom(const vector<VarOrRVar> &free_vars, int tag) {`
`124`		`- extra_args.emplace_back(tag);`
`125`	`130`	`for (const VarOrRVar &v : free_vars) {`
`126`	`131`	`if (v.is_rvar) {`
`127`	`132`	`extra_args.push_back(v.rvar);`
`128`	`133`	`} else {`
`129`	`134`	`extra_args.push_back(v.var);`
`130`	`135`	`}`
`131`	`136`	`}`
	`137`	`+ extra_args.emplace_back(tag);`
`132`	`138`	`}`
`133`	`139`	`};`
`134`	`140`