Merge pull request #899 from jamesr66a/primitive_lib

[FX] Add primitive library example

Author: James Reed

fx/primitive_library.py

@@ -0,0 +1,164 @@
+import torch
+import torch.fx
+"""
+In this example we are going do define a library of
+"composite" operations. Composite operations are those
+that are defined as callable functions that are composed
+of several other operations in their implementation.
+
+Composite operations allow you to choose at what level
+of abstraction you want to interpret/manipulate the
+code. We show that we can provide a function to inline
+these functions as well as use a custom Tracer to auto-
+matically inline such functions.
+
+Composite operations can be useful for exposing higher-
+level context to a backend/transform while still
+maintaining the ability to examine things at a more
+fine-grained level.
+"""
+
+
+def sigmoid_lowp(x : torch.Tensor):
+    x = x.float()
+    x = x.sigmoid()
+    return x.half()
+
+# wrap() indicates that the passed-in function should always
+# be recorded as a call_function node rather than being traced
+# through. Later, we will see how we can:
+# a. Inline the implementation of such a function and
+# b. Define a tracer that automatically traces through such a function
+torch.fx.wrap(sigmoid_lowp)
+
+def add_lowp(a : torch.Tensor, b : torch.Tensor):
+    a, b = a.float(), b.float()
+    c = a + b
+    return c.half()
+
+torch.fx.wrap(add_lowp)
+
+
+# Let's see what happens when we symbolically trace through some code
+# that uses these functions
+
+class Foo(torch.nn.Module):
+    def forward(self, x, y):
+        x = sigmoid_lowp(x)
+        y = sigmoid_lowp(y)
+        return add_lowp(x, y)
+
+
+traced = torch.fx.symbolic_trace(Foo())
+print(traced.code)
+"""
+def forward(self, x, y):
+    sigmoid_lowp = __main___sigmoid_lowp(x);  x = None
+    sigmoid_lowp_1 = __main___sigmoid_lowp(y);  y = None
+    add_lowp = __main___add_lowp(sigmoid_lowp, sigmoid_lowp_1);  sigmoid_lowp = sigmoid_lowp_1 = None
+    return add_lowp
+"""
+
+# Notice that the calls to `sigmoid_lowp` and `add_lowp`
+# appear literally in the trace; they are not traced through
+
+
+# ***** Inlining calls *****
+# Now let's define a function that allows for inlining these calls
+# during graph manipulation.
+
+def inline_lowp_func(n : torch.fx.Node):
+    # If we find a call to a function in our "lowp" module, inline it
+    if n.op == 'call_function' and n.target.__module__ == inline_lowp_func.__module__:
+        # We want to insert the operations comprising the implementation of the
+        # function before the function itself. Then, we can swap the output value
+        # of the function call with the output value for its implementation nodes
+        with n.graph.inserting_before(n):
+            # We can inline code by using `fx.Proxy` instances.
+            # map_arg traverses all aggregate types and applies the given function
+            # to Node instances in the data structure. In this case, we are applying
+            # the fx.Proxy constructor.
+            proxy_args = torch.fx.node.map_arg(n.args, torch.fx.Proxy)
+            proxy_kwargs = torch.fx.node.map_arg(n.kwargs, torch.fx.Proxy)
+            # Call the function itself with proxy arguments. This will emit
+            # nodes in the graph corresponding to the operations in the im-
+            # plementation of the function
+            output_proxy = n.target(*proxy_args, **proxy_kwargs)
+            # Now replace the original node's uses with the output node of
+            # the implementation.
+            node.replace_all_uses_with(output_proxy.node)
+            # Delete the old node
+            node.graph.erase_node(node)
+
+for node in traced.graph.nodes:
+    if node.op == 'call_function' and node.target is sigmoid_lowp:
+        inline_lowp_func(node)
+
+# Don't forget to recompile after graph manipulation
+traced.recompile()
+
+print(traced.code)
+"""
+def forward(self, x, y):
+    float_1 = x.float();  x = None
+    sigmoid = float_1.sigmoid();  float_1 = None
+    half = sigmoid.half();  sigmoid = None
+    float_2 = y.float();  y = None
+    sigmoid_1 = float_2.sigmoid();  float_2 = None
+    half_1 = sigmoid_1.half();  sigmoid_1 = None
+    add_lowp = __main___add_lowp(half, half_1);  half = half_1 = None
+    return add_lowp
+"""
+
+# At this point, the implementation of `sigmoid_lowp` has been substituted
+# in for all of the calls to that function.
+
+# ***** Inlining calls during tracing *****
+# Now we are going to define a custom tracer that can selectively inline
+# calls to certain composite operations on-the-fly.
+
+# New instance of our module
+f = Foo()
+
+class InliningTracer(torch.fx.Tracer):
+    FNS_TO_INLINE = [add_lowp]
+
+    def create_node(self, kind, target, args, kwargs, name=None, type_expr=None):
+        if kind == 'call_function' and target in self.FNS_TO_INLINE:
+            # Trace through the implementation of the function rather than
+            # create a node
+            proxy_args = torch.fx.node.map_arg(args, torch.fx.Proxy)
+            proxy_kwargs = torch.fx.node.map_arg(kwargs, torch.fx.Proxy)
+            return target(*proxy_args, **proxy_kwargs).node
+        else:
+            return super().create_node(kind, target, args, kwargs, name, type_expr)
+
+
+tracer = InliningTracer()
+graph = tracer.trace(f)
+module = torch.fx.GraphModule(f, graph)
+print(module.code)
+"""
+def forward(self, x, y):
+    sigmoid_lowp = __main___sigmoid_lowp(x);  x = None
+    sigmoid_lowp_1 = __main___sigmoid_lowp(y);  y = None
+    float_1 = sigmoid_lowp.float();  sigmoid_lowp = None
+    float_2 = sigmoid_lowp_1.float();  sigmoid_lowp_1 = None
+    add = float_1 + float_2;  float_1 = float_2 = None
+    half = add.half();  add = None
+    return half
+"""
+
+# As you can see, the implementation for `add_lowp` has been
+# inlined in the course of tracing with our InliningTracer.
+# Such functionality can be used to, for example, implement
+# a backend that wants to see the lowered form of some operations
+# but the high-level form of another.
+
+# ***** Future direction *****
+#
+# We may define an API, such as `Tracer.is_leaf_function`, that
+# Tracer implementers can use to more easily specify the inlining
+# behavior implemented in InliningTracer. Such a method would return
+# True by default, but a Tracer can override it and return `False` for
+# functions the Tracer wants to be traced through.