[Autodiff] Deterministic gradient compute

src/te/autodiff/ad_simplify.cc

@@ -413,15 +413,17 @@ class FactorOutAtomicFormulasFunctor
     auto res_b = VisitExpr(op->b);
 
     // For the And case we return the union of the sets of atomic formulas
-    std::unordered_set<PrimExpr, StructuralHash, StructuralEqual> res_set;
-    res_set.reserve(res_a.atomic_formulas.size() + res_b.atomic_formulas.size());
+    std::unordered_set<PrimExpr, StructuralHash, StructuralEqual> res_a_set;
+    res_a_set.reserve(res_a.atomic_formulas.size());
     std::copy(res_a.atomic_formulas.begin(), res_a.atomic_formulas.end(),
-              std::inserter(res_set, res_set.end()));
-    std::copy(res_b.atomic_formulas.begin(), res_b.atomic_formulas.end(),
-              std::inserter(res_set, res_set.end()));
-
-    std::vector<PrimExpr> res{res_set.begin(), res_set.end()};
+              std::inserter(res_a_set, res_a_set.end()));
 
+    std::vector<PrimExpr> res = res_a.atomic_formulas;
+    for (const auto& e : res_b.atomic_formulas) {
+      if (res_a_set.find(e) == res_a_set.end()) {
+        res.emplace_back(e);
+      }
+    }
     // And the residuals are combined with &&
     return {res, res_a.rest && res_b.rest};
   }
@@ -443,32 +445,34 @@ class FactorOutAtomicFormulasFunctor
 
     // For the Or case we intersect the sets of atomic formulas
     std::unordered_set<PrimExpr, StructuralHash, StructuralEqual> res_set;
+    std::vector<PrimExpr> res;
     res_set.reserve(std::min(res_a.atomic_formulas.size(), res_b.atomic_formulas.size()));
-    for (const auto& res_b_formula : res_b_set) {
+    res.reserve(std::min(res_a.atomic_formulas.size(), res_b.atomic_formulas.size()));
+    for (const auto& res_b_formula : res_b.atomic_formulas) {
       if (res_a_set.count(res_b_formula)) {
         res_set.insert(res_b_formula);
+        res.push_back(res_b_formula);
       }
     }
 
     // Computing the residual is more complex: we have to compute the sets of atomic formulas
     // which are left behind, and then combine them with the residuals into the new residual.
     std::vector<PrimExpr> new_cond_a;
     new_cond_a.reserve(res_a.atomic_formulas.size() - res_set.size());
-    for (const auto& formula : res_a_set) {
+    for (const auto& formula : res_a.atomic_formulas) {
       if (!res_set.count(formula)) new_cond_a.emplace_back(formula);
     }
 
     std::vector<PrimExpr> new_cond_b;
     new_cond_b.reserve(res_b.atomic_formulas.size() - res_set.size());
-    for (const auto& formula : res_b_set) {
+    for (const auto& formula : res_b.atomic_formulas) {
       if (!res_set.count(formula)) new_cond_b.emplace_back(formula);
     }
 
     res_a.atomic_formulas = std::move(new_cond_a);
     res_b.atomic_formulas = std::move(new_cond_b);
 
     PrimExpr new_rest = res_a.to_expr() || res_b.to_expr();
-    std::vector<PrimExpr> res{res_set.begin(), res_set.end()};
 
     return {res, new_rest};
   }

tests/python/unittest/test_te_autodiff.py

@@ -343,7 +343,23 @@ def test_reduction_init():
     check_grad(B, A0)
 
 
+def test_stable():
+    X = te.placeholder((32, 512, 16, 16), name="X")
+    W = te.placeholder((1024, 512, 1, 1), name="W")
+    strides, padding, dilation = 2, 0, 1
+    R = topi.nn.conv2d(X, W, strides, padding, dilation)
+    ones = topi.full_like(R, 1.0)
+    grads = te.gradient(R, [X], head=ones)
+    dag = tvm.auto_scheduler.ComputeDAG(grads)
+    repeat = 100
+    for i in range(repeat):
+        grads = te.gradient(R, [X], head=ones)
+        new_dag = tvm.auto_scheduler.ComputeDAG(grads)
+        assert str(dag) == str(new_dag)
+
+
 if __name__ == "__main__":
     test_basic_operation()
     test_topi()
     test_stride_dilation()
+    test_stable()