From Python to C/C++: a PyTorch walk-through

import torch
from torch.autograd import Variable

a = Variable(torch.FloatTensor([2]), requires_grad=True) # Parameters
b = Variable(torch.FloatTensor([3])) # Data
loss = torch.mm(a, b).sum()

loss.backward()

    class Function(object):
    # ...
    def _do_backward(self, grad_output):
        grad_input = self.backward(grad_output)
        # ... some trivial post-processing including `hook` and pack results
        return grad_input
    def backward(self, grad_output):
        raise NotImplementedError

class ExecutionEngine(object):

    def _compute_dependencies(self, function):
        """tranverse the computation graph to get the dependencies
        
        BFS tranverse the function starting from final output. The
        dependencies is a collection of counters.
        """
        # compute the dependencies
        return dependencies

    def _free_backward_dependency(self, dependencies, prev_fn, fn, arg_id):
        """Update the dependencies after backward one function
        
        Return:
            output_idx: the position of the arg in the outputs of prev_fn
        """
        # Update dependencies and return the position
        return output_idx


    def _is_ready_for_backward(self, dependencies, function):
        """Check if the node function is ready
        
        The ready status is determined by the in-degree of the node function
        a.k.a the dependencies[function][output] are all 0s.
        """
        for deps in dependencies[function]:
            if len(deps) > 0:
                return False
        return True

    def run_backward(self, variable, grad):
        """The core part of backward engine
        
        It calls the backward method of Functions and dispatches the gradients
        to the correct previous functions. The method is returned until all
        functions are `backward`ed successfully
        """
                                
        # set up the starting point, the grad is 1 without explicitly 
        # setting when calling `Variable.backward()`
        ready = [(variable.creator, (grad,))] # Functions ready for BP
        not_ready = {} # Functions whose gradients are not accumulated properly

        dependencies = self._compute_dependencies(variable.creator)

        while len(ready) > 0:
            fn, grad = ready.pop()
            # TODO: double-buffering
            grad_input = fn._do_backward(*grad)
                                    
            # Update the dependencies of all the input Variables of function fn
            # arg_id is used for position looking-up.
            for (prev_fn, arg_id), d_prev_fn in zip(fn.previous_functions, grad_input):
                                        
                # skipping useless gradients computing
                if not prev_fn.requires_grad:
                    assert d_prev_fn is None
                    continue
                output_nr = self._free_backward_dependency(dependencies, prev_fn, fn, arg_id)
                is_ready = self._is_ready_for_backward(dependencies, prev_fn)

                # the following codes are a little bit messy as
                # the two branches partially share a same underlying logic
                # , accumluating the current gradient to existing one.
                                            
                # if the `perv_fn` is ready for backward, then move
                # it from `not_ready` to `ready`
                if is_ready:
                    if prev_fn in not_ready:
                        prev_grad = not_ready[prev_fn]
                        if not prev_grad[output_nr]:
                            prev_grad[output_nr] = d_prev_fn
                        else:
                            prev_grad[output_nr].add_(d_prev_fn)
                        del not_ready[prev_fn]
                    else:
                        # The `prev_fn` is ready when the first seen,
                        # there must be only one output in `prev_fn`
                        assert output_nr == 0
                        prev_grad = (d_prev_fn,)
                    ready.append((prev_fn, prev_grad))
                else:
                    if prev_fn in not_ready:
                        prev_grad = not_ready[prev_fn]
                    else:
                        prev_grad = [None for _ in prev_fn.output_ids]

                    if not prev_grad[output_nr]:
                        prev_grad[output_nr] = d_prev_fn
                    else:
                        prev_grad[output_nr].add_(d_prev_fn)

                    not_ready[prev_fn] = prev_grad

class ExecutionEngine(object):
    def __init__(self):
        pass

    def _compute_dependencies(self, function):
        """tranverse the computation graph to get the dependencies
        
        BFS tranverse the function starting from final output. The
        dependencies is a collection of counters.
        """
        dependencies = {}
        seen = {function}
        queue = [function]
        while len(queue) > 0:
            fn = queue.pop()
            for prev_fn, arg_id in fn.previous_functions:
                if prev_fn not in dependencies:
                    dependencies[prev_fn] = [Counter() for _ in prev_fn.output_ids]
                output_idx = prev_fn.output_ids[arg_id]
                    
                # I think there's no need to store the counter for each current function,
                # a simple counter for each prev_fn should be enough.
                dependencies[prev_fn][output_idx][fn] += 1
                if prev_fn not in seen:
                    queue.append(prev_fn)
                    seen.add(prev_fn)
        return dependencies

    def _free_backward_dependency(self, dependencies, prev_fn, fn, arg_id):
        """Update the dependencies after backward one function
        
        Return:
            output_idx: the position of the arg in the outputs of prev_fn
        """
        deps = dependencies[prev_fn]
        output_idx = prev_fn.output_ids[arg_id]
        output_deps = deps[output_idx]
        output_deps[fn] -= 1
        if output_deps[fn] == 0:
            del output_deps[fn]
        return output_idx


    def _is_ready_for_backward(self, dependencies, function):
        """Check if the node function is ready
        
        The ready status is determined by the in-degree of the node function
        a.k.a the dependencies[function][output] are all 0s.
        """
        for deps in dependencies[function]:
            if len(deps) > 0:
                return False
        return True