Reuse KernelSpec for FusionGroups with equivalent graphs (#14541)
authorRichard Zou <zou3519@gmail.com>
Thu, 13 Dec 2018 15:51:08 +0000 (07:51 -0800)
committerFacebook Github Bot <facebook-github-bot@users.noreply.github.com>
Thu, 13 Dec 2018 15:54:35 +0000 (07:54 -0800)
commitb14d6d730a7fac0d6d31be600878d1a7ca4af217
tree31031c04a9b09518a1cd1adc569199f7a8608f3e
parentaa022313cb694546cca7fcb0123a818bf2d668b9
Reuse KernelSpec for FusionGroups with equivalent graphs (#14541)

Summary:
Before this PR, loop unrolling + the graph fuser was creating multiple
FusionGroups with the same bodies (with different variable names) for
JIT LSTMs. Each FusionGroup got registered to a separate fusion key;
each key resulted in a different compilation for the same
specializations.

This PR makes it so that when registering FusionGroups with the fusion
compiler, the compiler first checks the KernelSpec cache to see if the
FusionGroup's graph exists already. If it does, then return the
corresponding KernelSpec's key to share compiled kernels.

In addition, graphs in the KernelSpec cache are canonicalized before
being cached. I added a flag to the canonicalize pass to remove unique
names of values.

This shortens the compile time for a JIT LSTM (seq_len of 100, loop
unroll factor of 8) from 5.3s to 2.3s. Most of this compile time is
running the graph fuser and/or fusion compiler; while this PR
makes it so that there is only one unique kernel in the forward pass,
there are a lot of different kernels (6) in the backward pass
(after loop unrolling) that should be investigated.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/14541

Differential Revision: D13324487

Pulled By: zou3519

fbshipit-source-id: b841d82ed35a959b5cfc72db033bf5a7b42cc4fb
test/cpp/jit/gtest.cpp
test/cpp/jit/tests.h
test/test_jit.py
torch/csrc/jit/fuser/compiler.cpp
torch/csrc/jit/fuser/compiler.h
torch/csrc/jit/fuser/kernel_cache.cpp
torch/csrc/jit/fuser/kernel_cache.h
torch/csrc/jit/init.cpp
torch/csrc/jit/passes/canonicalize.cpp
torch/csrc/jit/passes/canonicalize.h