Add a cheap hash to Expr nodes to speed up IREquality

src/Expr.cpp

@@ -1,3 +1,5 @@
+#include <functional>
+
 #include "Expr.h"
 #include "IROperator.h"  // for lossless_cast()
 
@@ -35,6 +37,12 @@ const IntImm *IntImm::make(Type t, int64_t value) {
     IntImm *node = new IntImm;
     node->type = t;
     node->value = value;
+    // Small values are extremely common, so a hash that just slices up the
+    // bits of the value (like combine_hash below) would put all the entropy
+    // for those in the low bits, which get discarded by set_hash. Multiply
+    // by a large odd constant and keep the high bits instead, which mixes
+    // in the low bits of the value even when the value itself is small.
+    node->set_hash((uint32_t)((((uint64_t)value) * 0x9e3779b97f4a7c15ULL) >> 32));
     return node;
 }
 
@@ -51,6 +59,9 @@ const UIntImm *UIntImm::make(Type t, uint64_t value) {
     UIntImm *node = new UIntImm;
     node->type = t;
     node->value = value;
+    // See the comment in IntImm::make about why we multiply rather than
+    // just slicing up the bits of the value.
+    node->set_hash((uint32_t)((value * 0x9e3779b97f4a7c15ULL) >> 32));
     return node;
 }
 
@@ -77,13 +88,15 @@ const FloatImm *FloatImm::make(Type t, double value) {
         internal_error << "FloatImm must be 16, 32, or 64-bit\n";
     }
 
+    node->set_hash((uint32_t)std::hash<double>{}(node->value));
     return node;
 }
 
 const StringImm *StringImm::make(const std::string &val) {
     StringImm *node = new StringImm;
     node->type = type_of<const char *>();
     node->value = val;
+    node->set_hash((uint32_t)std::hash<std::string>{}(val));
     return node;
 }
 

src/Expr.h

@@ -104,7 +104,7 @@ struct IRNode {
      */
     virtual void accept(IRVisitor *v) const = 0;
     IRNode(IRNodeType t)
-        : node_type(t) {
+        : hash((uint32_t)t) {
     }
     virtual ~IRNode() = default;
 
@@ -115,17 +115,34 @@ struct IRNode {
      */
     mutable RefCount ref_count;
 
-    /** Each IR node subclass has a unique identifier. We can compare
-     * these values to do runtime type identification. We don't
-     * compile with rtti because that injects run-time type
-     * identification stuff everywhere (and often breaks when linking
-     * external libraries compiled without it), and we only want it
-     * for IR nodes. One might want to put this value in the vtable,
-     * but that adds another level of indirection, and for Exprs we
-     * have 32 free bits in between the ref count and the Type
-     * anyway, so this doesn't increase the memory footprint of an IR node.
-     */
-    IRNodeType node_type;
+    /** Each IR node subclass has a unique identifier. We can compare these
+     * values to do runtime type identification. We don't compile with rtti
+     * because that injects run-time type identification stuff everywhere (and
+     * often breaks when linking external libraries compiled without it), and we
+     * only want it for IR nodes. One might want to put this value in the
+     * vtable, but that adds another level of indirection, and for Exprs we have
+     * 32 free bits in between the ref count and the Type field anyway. We use
+     * the first 8 to store the node type, and the next 24 as a hash of the
+     * children of the node, to make syntactic comparisons faster. */
+    union {
+        IRNodeType node_type;
+        uint32_t hash;
+    };
+
+    /** Set hash from a combined hash of this node's arguments (see
+     * combine_hash below), keeping the node type intact. The low bits of a
+     * multiply-add hash are of poor quality, so we discard them (rather
+     * than shifting them up) in favor of the node type. Which end of the
+     * word the node type landed in when we wrote it via the node_type
+     * member of the union depends on the endianness of the machine. */
+    HALIDE_ALWAYS_INLINE
+    void set_hash(uint32_t args_hash) {
+#if defined(__BYTE_ORDER__) && __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
+        hash = (args_hash >> 8) | ((uint32_t)node_type << 24);
+#else
+        hash = (args_hash & 0xffffff00u) | (uint32_t)node_type;
+#endif
+    }
 };
 
 template<>
@@ -163,6 +180,22 @@ struct BaseExprNode : public IRNode {
     Type type;
 };
 
+/** Combine one or more child hashes (or plain uint32_t fields) into a
+ * running hash, for use in the make() methods of Expr nodes below. Pass
+ * the result to IRNode::set_hash to fold in the node type and get the
+ * final hash - see the make() methods below for examples. */
+// @{
+HALIDE_ALWAYS_INLINE
+uint32_t combine_hash(uint32_t hash, uint32_t child_hash) {
+    return hash * 2654435761u + child_hash;
+}
+
+template<typename... Rest>
+HALIDE_ALWAYS_INLINE uint32_t combine_hash(uint32_t hash, uint32_t child_hash, Rest... rest) {
+    return combine_hash(combine_hash(hash, child_hash), rest...);
+}
+// @}
+
 /** We use the "curiously recurring template pattern" to avoid
    duplicated code in the IR Nodes. These classes live between the
    abstract base classes and the actual IR Nodes in the
@@ -342,6 +375,12 @@ struct Expr : public Internal::IRHandle {
     Type type() const {
         return get()->type;
     }
+
+    /** Get the cheap hash of this expression node. See IRNode::hash. */
+    HALIDE_ALWAYS_INLINE
+    uint32_t hash() const {
+        return get()->hash;
+    }
 };
 
 /** This lets you use an Expr as a key in a map of the form

src/IR.cpp

@@ -4,6 +4,7 @@
 #include "IROperator.h"
 #include "IRPrinter.h"
 #include "IRVisitor.h"
+#include <functional>
 #include <numeric>
 #include <utility>
 
@@ -44,6 +45,7 @@ Expr Cast::make(Type t, Expr v) {
 
     Cast *node = new Cast;
     node->type = t;
+    node->set_hash(combine_hash(v.hash(), t.hash()));
     node->value = std::move(v);
     return node;
 }
@@ -60,6 +62,7 @@ Expr Reinterpret::make(Type t, Expr v) {
 
     Reinterpret *node = new Reinterpret;
     node->type = t;
+    node->set_hash(combine_hash(v.hash(), t.hash()));
     node->value = std::move(v);
     return node;
 }
@@ -71,6 +74,7 @@ Expr Add::make(Expr a, Expr b) {
 
     Add *node = new Add;
     node->type = a.type();
+    node->set_hash(combine_hash(a.hash(), b.hash()));
     node->a = std::move(a);
     node->b = std::move(b);
     return node;
@@ -83,6 +87,7 @@ Expr Sub::make(Expr a, Expr b) {
 
     Sub *node = new Sub;
     node->type = a.type();
+    node->set_hash(combine_hash(a.hash(), b.hash()));
     node->a = std::move(a);
     node->b = std::move(b);
     return node;
@@ -95,6 +100,7 @@ Expr Mul::make(Expr a, Expr b) {
 
     Mul *node = new Mul;
     node->type = a.type();
+    node->set_hash(combine_hash(a.hash(), b.hash()));
     node->a = std::move(a);
     node->b = std::move(b);
     return node;
@@ -107,6 +113,7 @@ Expr Div::make(Expr a, Expr b) {
 
     Div *node = new Div;
     node->type = a.type();
+    node->set_hash(combine_hash(a.hash(), b.hash()));
     node->a = std::move(a);
     node->b = std::move(b);
     return node;
@@ -119,6 +126,7 @@ Expr Mod::make(Expr a, Expr b) {
 
     Mod *node = new Mod;
     node->type = a.type();
+    node->set_hash(combine_hash(a.hash(), b.hash()));
     node->a = std::move(a);
     node->b = std::move(b);
     return node;
@@ -131,6 +139,7 @@ Expr Min::make(Expr a, Expr b) {
 
     Min *node = new Min;
     node->type = a.type();
+    node->set_hash(combine_hash(a.hash(), b.hash()));
     node->a = std::move(a);
     node->b = std::move(b);
     return node;
@@ -143,6 +152,7 @@ Expr Max::make(Expr a, Expr b) {
 
     Max *node = new Max;
     node->type = a.type();
+    node->set_hash(combine_hash(a.hash(), b.hash()));
     node->a = std::move(a);
     node->b = std::move(b);
     return node;
@@ -155,6 +165,7 @@ Expr EQ::make(Expr a, Expr b) {
 
     EQ *node = new EQ;
     node->type = Bool(a.type().lanes());
+    node->set_hash(combine_hash(a.hash(), b.hash()));
     node->a = std::move(a);
     node->b = std::move(b);
     return node;
@@ -167,6 +178,7 @@ Expr NE::make(Expr a, Expr b) {
 
     NE *node = new NE;
     node->type = Bool(a.type().lanes());
+    node->set_hash(combine_hash(a.hash(), b.hash()));
     node->a = std::move(a);
     node->b = std::move(b);
     return node;
@@ -179,6 +191,7 @@ Expr LT::make(Expr a, Expr b) {
 
     LT *node = new LT;
     node->type = Bool(a.type().lanes());
+    node->set_hash(combine_hash(a.hash(), b.hash()));
     node->a = std::move(a);
     node->b = std::move(b);
     return node;
@@ -191,6 +204,7 @@ Expr LE::make(Expr a, Expr b) {
 
     LE *node = new LE;
     node->type = Bool(a.type().lanes());
+    node->set_hash(combine_hash(a.hash(), b.hash()));
     node->a = std::move(a);
     node->b = std::move(b);
     return node;
@@ -203,6 +217,7 @@ Expr GT::make(Expr a, Expr b) {
 
     GT *node = new GT;
     node->type = Bool(a.type().lanes());
+    node->set_hash(combine_hash(a.hash(), b.hash()));
     node->a = std::move(a);
     node->b = std::move(b);
     return node;
@@ -215,6 +230,7 @@ Expr GE::make(Expr a, Expr b) {
 
     GE *node = new GE;
     node->type = Bool(a.type().lanes());
+    node->set_hash(combine_hash(a.hash(), b.hash()));
     node->a = std::move(a);
     node->b = std::move(b);
     return node;
@@ -229,6 +245,7 @@ Expr And::make(Expr a, Expr b) {
 
     And *node = new And;
     node->type = Bool(a.type().lanes());
+    node->set_hash(combine_hash(a.hash(), b.hash()));
     node->a = std::move(a);
     node->b = std::move(b);
     return node;
@@ -243,6 +260,7 @@ Expr Or::make(Expr a, Expr b) {
 
     Or *node = new Or;
     node->type = Bool(a.type().lanes());
+    node->set_hash(combine_hash(a.hash(), b.hash()));
     node->a = std::move(a);
     node->b = std::move(b);
     return node;
@@ -254,6 +272,7 @@ Expr Not::make(Expr a) {
 
     Not *node = new Not;
     node->type = Bool(a.type().lanes());
+    node->set_hash(combine_hash(a.hash(), 0));
     node->a = std::move(a);
     return node;
 }
@@ -269,6 +288,8 @@ Expr Select::make(Expr condition, Expr true_value, Expr false_value) {
 
     Select *node = new Select;
     node->type = true_value.type();
+    node->set_hash(combine_hash(condition.hash(),
+                                true_value.hash(), false_value.hash()));
     node->condition = std::move(condition);
     node->true_value = std::move(true_value);
     node->false_value = std::move(false_value);
@@ -285,6 +306,8 @@ Expr Load::make(Type type, const std::string &name, Expr index, Buffer<> image,
     Load *node = new Load;
     node->type = type;
     node->name = name;
+    node->set_hash(combine_hash((uint32_t)std::hash<std::string>{}(name),
+                                index.hash(), predicate.hash()));
     node->predicate = std::move(predicate);
     node->index = std::move(index);
     node->image = std::move(image);
@@ -320,6 +343,8 @@ Expr Ramp::make(Expr base, Expr stride, int lanes) {
 
     Ramp *node = new Ramp;
     node->type = base.type().with_lanes(lanes * base.type().lanes());
+    node->set_hash(combine_hash((uint32_t)lanes,
+                                base.hash(), stride.hash()));
     node->base = std::move(base);
     node->stride = std::move(stride);
     node->lanes = lanes;
@@ -332,6 +357,7 @@ Expr Broadcast::make(Expr value, int lanes) {
 
     Broadcast *node = new Broadcast;
     node->type = value.type().with_lanes(lanes * value.type().lanes());
+    node->set_hash(combine_hash((uint32_t)lanes, value.hash()));
     node->value = std::move(value);
     node->lanes = lanes;
     return node;
@@ -344,6 +370,8 @@ Expr Let::make(const std::string &name, Expr value, Expr body) {
     Let *node = new Let;
     node->type = body.type();
     node->name = name;
+    node->set_hash(combine_hash((uint32_t)std::hash<std::string>{}(name),
+                                value.hash(), body.hash()));
     node->value = std::move(value);
     node->body = std::move(body);
     return node;
@@ -974,6 +1002,12 @@ Expr Call::make(Type type, const std::string &name, const std::vector<Expr> &arg
     Call *node = new Call;
     node->type = type;
     node->name = name;
+    uint32_t h = combine_hash((uint32_t)std::hash<std::string>{}(name),
+                              (uint32_t)call_type, (uint32_t)value_index);
+    for (const auto &arg : args) {
+        h = combine_hash(h, arg.hash());
+    }
+    node->set_hash(h);
     node->args = args;
     node->call_type = call_type;
     node->func = std::move(func);
@@ -995,6 +1029,7 @@ Expr Variable::make(Type type, const std::string &name, Buffer<> image, Paramete
     Variable *node = new Variable;
     node->type = type;
     node->name = name;
+    node->set_hash((uint32_t)std::hash<std::string>{}(name));
     node->image = std::move(image);
     node->param = std::move(param);
     node->reduction_domain = std::move(reduction_domain);
@@ -1017,6 +1052,14 @@ Expr Shuffle::make(const std::vector<Expr> &vectors,
 
     Shuffle *node = new Shuffle;
     node->type = element_ty.with_lanes((int)indices.size());
+    uint32_t h = 0;
+    for (int i : indices) {
+        h = combine_hash(h, (uint32_t)i);
+    }
+    for (const auto &v : vectors) {
+        h = combine_hash(h, v.hash());
+    }
+    node->set_hash(h);
     node->vectors = vectors;
     node->indices = indices;
     return node;
@@ -1259,6 +1302,7 @@ Expr VectorReduce::make(VectorReduce::Operator op,
         << lanes << " " << vec.type().lanes() << "\n";
     VectorReduce *node = new VectorReduce;
     node->type = vec.type().with_lanes(lanes);
+    node->set_hash(combine_hash((uint32_t)op, vec.hash()));
     node->op = op;
     node->value = std::move(vec);
     return node;