[functorch] fix distance heuristic

Author: Chillee

functorch/functorch/_src/decompositions.py

@@ -12,12 +12,21 @@
 # Decompositions have been ported to torch._decomp inside of PyTorch core. The only decompositions here are temporary or hacks. Please submit your contributions to PyTorch core!
 
 
-@register_decomposition(aten.trace.default)
+def maybe_register_decomposition(op):
+    def decorator(f):
+        try:
+            return register_decomposition(op)(f)
+        except Exception:
+            return f
+    return decorator
+
+
+@maybe_register_decomposition(aten.trace.default)
 def trace(self: Tensor) -> Tensor:
     return torch.sum(torch.diag(self))
 
 
-@register_decomposition(aten.log_sigmoid_forward)
+@maybe_register_decomposition(aten.log_sigmoid_forward.default)
 def log_sigmoid_forward(self: Tensor) -> Tuple[Tensor, Tensor]:
     min = torch.minimum(self.new_zeros(()), self)
     z = torch.exp(-torch.abs(self))

functorch/functorch/_src/partitioners.py

@@ -251,21 +251,23 @@ def classify_nodes(joint_module):
         return required_fw_nodes, required_bw_nodes, unclaimed_nodes
 
     required_fw_nodes, required_bw_nodes, unclaimed_nodes = classify_nodes(joint_module)
-    cache = {}
 
-    def dist_from_fw(node):
-        if node in cache:
-            return cache[node]
+    def dist_from_bw(node):
+        print(node)
         if node not in required_fw_nodes:
             return 0
         dist = int(1e9)
         for n in node.users:
-            dist = min(dist_from_fw(n) + 1, dist)
-        cache[node] = dist
+            dist = min(dist_from_bw(n) + 1, dist)
         return dist
 
-    for node in joint_module.graph.nodes:
-        node.dist_from_fw = dist_from_fw(node)
+    for node in reversed(joint_module.graph.nodes):
+        if node not in required_fw_nodes:
+            node.dist_from_bw = 0
+        else:
+            node.dist_from_bw = int(1e9)
+            for n in node.users:
+                node.dist_from_bw = min(node.dist_from_bw, n.dist_from_bw + 1)
 
     aten = torch.ops.aten
 
@@ -325,7 +327,7 @@ def get_node_weight(node):
         mem_sz = _size_of(node.meta['tensor_meta'])
 
         # Heuristic to bias towards nodes closer to the backwards pass
-        mem_sz = int(mem_sz + node.dist_from_fw)
+        mem_sz = int(mem_sz + node.dist_from_bw)
 
         if is_materialized(node):
             return mem_sz