AITemplate Update fuse_split_linear

Update fuse_split_linear_add

Open JiaJiunn opened this issue 1 year ago • 5 comments

Summary: Currently, fuse_split_linear_add only supports cases when the split op kwarg uses a slice. This diff extends the fusion to support cases when the split op kwarg uses ints. The graph change can seen as follows:

Before optimization: {F933055696}
After optimization: {F933058824}

In particular, the optimization targets nodes split_5 and split_7 specifically.

Additionally, modified common_fx2ait to use logger instead of print, as the print statement causes buck tests to crash.

Differential Revision: D44359231