[1.x][FEATURE] CUDA graphs support

[ptrendx]

Description

CUDA graphs is a feature of CUDA 10, which enables lowering CPU overhead by bundling multiple kernel launches together.

The main limitation of CUDA graphs is that they do require the graph to be static - no parameters to the kernels can change, otherwise the graph needs to be recaptured. That is why this feature is currently only enabled for the symbolic models and Gluon models hybridized with hybridize(static_alloc=True, static_shape=True).

The feature is not enabled by default and requires environment variable MXNET_ENABLE_CUDA_GRAPHS to be set. In order to not capture the operations, the execution of which may change during the course of the training job, stateful operators and operators relying on resources other than workspace are not included in the graph.

Since the feature lowers the CPU overhead of the execution, the impact is most visible in the inference or scale-out workloads with small batch size. Let us consider following script comparing fp16, batch size 1 inference of RN50v2 model from GluonCV with and without graphs:

import mxnet as mx
import gluoncv as gcv
from gluoncv.model_zoo import get_model
import time
import os

net = get_model('ResNet50_V2')
net2 = get_model('ResNet50_V2')

net.initialize(ctx=mx.gpu())
net2.initialize(ctx=mx.gpu())

net.cast('float16')
net2.cast('float16')

net.hybridize(static_alloc=True, static_shape=True)
net2.hybridize(static_alloc=True, static_shape=True)

img = mx.random.uniform(shape=(1, 3, 224, 224), ctx=mx.gpu(), dtype='float16')

os.environ["MXNET_ENABLE_CUDA_GRAPHS"] = "0"

for _ in range(10):
    o = net(img)

mx.nd.waitall()

s = time.time()
for _ in range(1000):
    o = net(img)

mx.nd.waitall()
e = time.time()
print("No graphs: ", e - s)
mx.nd.waitall()

os.environ["MXNET_ENABLE_CUDA_GRAPHS"] = "1"

for _ in range(10):
    o = net2(img)

mx.nd.waitall()

s = time.time()
for _ in range(1000):
    o = net2(img)

mx.nd.waitall()
e = time.time()
print("With graphs: ", e - s)

The result obtained on V100 16GB:

[20:33:46] ../src/operator/nn/./cudnn/./cudnn_algoreg-inl.h:97: Running performance tests to find the best convolution algorithm, this can take a while... (set the environment variable MXNET_CUDNN_AUTOTUNE_DEFAULT to 0 to disable)
No graphs:  2.8153152465820312
With graphs:  2.3230245113372803

so over 17% increase in performance. The same script but with 128 batch size gives much smaller improvement:

[20:38:34] ../src/operator/nn/./cudnn/./cudnn_algoreg-inl.h:97: Running performance tests to find the best convolution algorithm, this can take a while... (set the environment variable MXNET_CUDNN_AUTOTUNE_DEFAULT to 0 to disable)
No graphs:  47.31952977180481
With graphs:  47.152849197387695

so 0.3%.

Checklist

Essentials

• PR's title starts with a category (e.g. [BUGFIX], [MODEL], [TUTORIAL], [FEATURE], [DOC], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented

[mxnet-bot]

Hey @ptrendx , Thanks for submitting the PR
All tests are already queued to run once. If tests fail, you can trigger one or more tests again with the following commands:

• To trigger all jobs: @mxnet-bot run ci [all]
• To trigger specific jobs: @mxnet-bot run ci [job1, job2]

---

CI supported jobs: [clang, centos-cpu, windows-gpu, windows-cpu, unix-cpu, website, miscellaneous, sanity, edge, unix-gpu, centos-gpu]

---

Note:
Only following 3 categories can trigger CI :PR Author, MXNet Committer, Jenkins Admin.
All CI tests must pass before the PR can be merged.

[ptrendx]

Forgot to add in the description that most of the work here was done by @DickJC123.

[samskalicky]

istead of setting false everywhere, can we just check if hasAttr("FIsCUDAGraphsCompatible") so that by default its false?

[ptrendx]

Well, the "everywhere" here is actually way less than if we went the other way around (only a bunch of operators with FCompute has synchronization that is not allowed under graphs.

[ptrendx]

And the long term goal here is actually to make those excluded operators compatible too.

[samskalicky]

Do something like this?
https://github.com/apache/incubator-mxnet/blob/72eff9b66ecc683c3e7f9ad2c0ba69efa8dd423b/src/imperative/imperative_utils.h#L83-L84

[ptrendx]

Not sure I understand - I do want to have the function call per Op here, not just a simple true/false.

[samskalicky]

Whats the point of registering a lambda function that just returns false?

[ptrendx]

I guess look here: https://github.com/apache/incubator-mxnet/pull/19142/files#diff-789523bf443903e74acfa010a5d6b572R33-R37 - this is for dropout, which uses random resource for training (and thus is excluded), but is just a passthrough for inference (and so we want to include it there).

[samskalicky]

ya that one is fine, but "_npi_eig" is always false...

[ptrendx]

But then I would need to add that function for almost all operators for it to return true...

[ptrendx]

We went with the default to be include the operator, but instead have the functionality itself be non-default.

[samskalicky]

Are the majority of ops supported to be included in cuda graphs? If a new user comes in to write an op, do they need to be aware of how to handle cuda graph support?

[ptrendx]

For the first question: yes - nearly all of the FCompute functions that do not use random resource are compatible with graphs.
For the second question - for FCompute operators to not be compatible you need to do synchronization inside the operator - either via stream synchronize or allocation. You generally do not want to do either one of those as it really hurts performance (and the operators that in this PR I marked incompatible in this PR do just that). If you just launch a kernel (or multiple), which is the case for vast majority of the operators, then you are good and do not even need to think about graphs - it will just work.

I'm still exploring the ways of automatically testing newly added operators in order for the feature to be able to be on by default, but I do not consider this the scope of this PR, as v1.x branch is not really supposed to get many more operators (I will do that in the PR to master). Generally this would involve testing operators with hybridize(static_alloc=True, static_shape=True) (which generally should be tested much more as right now testing of this functionality is really limited, even though it is widely used).

[samskalicky]

I think we should be ok since its disabled by default, and only enabled explicitly by setting the env var.

[DickJC123]

As @ptrendx mentioned, I coded much of this facility, as embodied in the cuda_graphs.h file. I have reviewed that file carefully and found it to be true to the version we have been using in NVIDIA's container version of MXNET. Thanks @ptrendx for adding the FIsCUDAGraphsCompatible attribute, as is needed for a robust CUDA Graphs support, and other improvements.

LGTM.