[Schedule] Add .fuse() primitive

[chhzh123]

Description

This PR adds a new primitive called .fuse(subgraph, compiler) for operator fusion. Currently we only support pattern-based vertical operator fusion using TorchScript as the backend compiler. A simple example is shown below.

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 32, 3)

    def forward(self, x):
        x = self.conv(x)
        x = F.relu(x)
        x = x + 1
        return x

mod = Model()
sch = slapo.create_schedule(mod)

# Define a function pattern whose ops will be fused
def pattern(x: torch.Tensor):
    x = F.relu(x)
    x = x + 1
    return x

# Find the subgraph in the original module that matches the specified pattern
subgraph = sch.find(pattern)
# Fuse the subgraph
sch.fuse(subgraph, compiler="TorchScript", name="FusedReLU")

The fusion performance needs to be tested. As I first create an identical torch.fx subgraph and pass it into TorchScript's scripting mode for optimization, it may cause performance issue if some of the operators are not recognized by TorchScript. Tracing mode can achieve the best performance but it is hard to leverage since we cannot always obtain example inputs for each subgraph. Scripting a user-defined function is also not a good approach since the backward pass is not captured by TorchScript. Therefore, only scripting an entire module is a good fit for our case, and we need to test the compatibility of torch.fx and TorchScript.

Checklist

• Support using op in torch.nn.functional to match modules in torch.nn. For example, F.relu and nn.ReLU should be treated as the same in pattern matching, since users cannot specify a module in a function pattern. ([Schedule] Refactor subgraph matching #35)
• Support subgraph matching with multiple input arguments. ([Schedule] Refactor subgraph matching #35)
• Add .decompose() primitive and support decoupling bias from nn.Linear
• Support generate a flattened traced graph ([Tracer] Add flatten argument to .trace() #29)
• Fuse the decomposed bias with consequential operators, e.g. bias+GeLU.

Future plan (Updated)

The following features will be added in separate PRs.

• Subgraph matching with multiple input arguments ([Schedule] Refactor subgraph matching #35)
• Extensible compiler backend (e.g., TVM)
• Rule-based subgraph matching ([Schedule] Refactor subgraph matching #35), probably needs more use cases and testings

[comaniac]

Seems like we could ask for sample inputs in kwargs if required by the compiler backend. For example: compiler="torchscript", example_inputs=...

[chhzh123]

It seems this PR requires several changes to the subgraph matching mechanism. I will open a new PR to make step 1 & 2 work first.

[comaniac]

Overall LGTM. Later we should apply .fuse to the examples, which currently replace an entire MLP for bias_new_gelu.

[comaniac]

The names of CallModule and call_module are confusing.

[chhzh123]

They are just used in different places. CallModule is used in a pattern class, while call_module can be used in a pattern function. Do you have any suggestions about the naming?

[comaniac]

Hmm base on your use cases, such as

    class LinearReLUPattern(slapo.Pattern):
        def __init__(self):
            super().__init__()
            self.fc = CallModule(r"fc?")
            self.relu = nn.ReLU()

Here you're actually constructing a module instead of calling it, so I guess you could use the name like ModulePattern or something like that.

[chhzh123]

It may also lead to confusion with the Pattern class.

[comaniac]

Well, to me it is a kind of pattern, especially you have a regex in its argument.

[chhzh123]

Yes, I agree it is a pattern, and ModulePattern is a proper name, but since we have the Pattern base class which is also a module, I'm afraid it will cause confusion between these two.

[chhzh123]

Later we should apply .fuse to the examples, which currently replace an entire MLP for bias_new_gelu.

Sure, I'll add it in the next PR.

[comaniac]

Thanks @chhzh123