{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "\n# Add Sub-Graph Rewrite Rule\n\nThis tutorial shows how to add a sub-graph rewrite rule in the graph optimization pipeline. Sub-graph rewriting is an\nimportant technique in graph optimization. It is used to replace a sub-graph with another sub-graph, which is usually\nmore efficient than the original one. For example, we can replace a sub-graph with two matrix multiplications sharing\nthe same input and one addition with a concatenation and a single matrix multiplication:\n\n.. figure:: /_static/img/subgraph-rewrite-example.svg\n :align: center\n :scale: 70%\n\n The sub-graph rewrite rule that fuses two matrix multiplications.\n\n.. seealso::\n :class: margin\n\n [TASO](https://dl.acm.org/doi/10.1145/3341301.3359630) systematically studies the sub-graph rewrite optimization\n for deep learning workloads.\n\nAfter the rewrite, the graph becomes more efficient as we only need to run a single kernel and the `fused` matrix\nmultiplication usually exposes more parallelism to utilize the underlying hardware. We can also fuse multiple\nconvolutions into a single one, or do other sub-graph rewrites.\n\n## Sub-graph rewrite in Hidet\n\nIn Hidet, we use a *sub-graph rewrite rule* to describe the rewrite. A sub-graph rewrite rule contains two parts:\n\n- **Sub-graph pattern**: a sub-graph pattern that we use to match the sub-graph in the graph. The pattern is a directed\n acyclic graph (DAG) where each node is an operator pattern and each edge is a tensor pattern. We only specify the\n operator type for each node and whether the (input) tensors are symbolic or concrete.\n\n- **Target sub-graph constructor**: when we find a sub-graph that matches the pattern, we use the constructor to\n construct the target sub-graph that replaces the matched sub-graph. When constructing the target sub-graph, we can\n use the matched tensors and nodes to further determine whether the rewrite is applicable. If applicable, the\n constructor returns the target sub-graph, otherwise it returns ``None``.\n\nIn the above example, the sub-graph pattern contains three input tensors, where x1 is a symbolic tensor and w1, w2 are\ntwo concrete tensors (i.e., we know the concrete values of w1 and w2). There are three operators in the pattern, where\nthe first two are matrix multiplications and the last one is an addition. The output tensor of the addition is the\noutput tensor of the pattern. When we find a sub-graph that matches the pattern, we use the constructor to construct\nthe target sub-graph and replace the matched sub-graph with the target sub-graph.\n\n

#### Note

**Difference between sub-graph rewrite and operator resolving**. Although\n `operator resolving ` can be conceptually considered as a special case of\n sub-graph rewrite, we use a different mechanism to implement them and their execution order is different. The operator\n resolving can be performed efficiently thus we can add arbitrary number of operator resolve rules. But the sub-graph\n rewrite is usually more expensive. Second, we run the sub-graph rewrite pass before the operator resolving pass, so\n that we can use the generic operators in the sub-graph patterns without worrying about the operator resolving.