Add static dispatch mode to reduce mobile code size by li-roy · Pull Request #22335 · pytorch/pytorch

li-roy · 2019-06-28T09:12:37Z

Stack from ghstack:

Add static dispatch mode to reduce mobile code size #22335 Add static dispatch mode to reduce mobile code size

As we discussed, this will allow the linker to remove unused operators automatically.

Differential Revision: D16048264

Add static dispatch mode to reduce mobile code size gh-metadata: pytorch pytorch 22335 gh/li-roy/36/head

smessmer

There's a few questions we should talk about before landing this, but generally looks good.

smessmer · 2019-06-28T20:31:35Z

.circleci/config.yml


+        # dispatch aten ops statically for mobile
+        if [[ ${BUILD_ENVIRONMENT} == *"android"* ]]; then
+          NAMED_FLAG="export USE_STATIC_DISPATCH=1"


Let's give it a more scary name. I don't want people to think "static dispatch sounds nice" and use this flag without knowing that it restricts them to CPU only, prevents them from overriding kernels, and has a few other restrictions.

Actually, thinking about it, maybe we should combine this flag with the flag we're planning for selecting which ops should be registered? Have a flag

PYTORCH_WHITELIST_ATEN_OPS_FOR_MOBILE="aten::conv,aten::add"

or

PYTORCH_ONLY_BUILD_OPS_FOR_MOBILE="aten::conv,aten::add"

and whenever this flag is present, you use static dispatch?

Yeah I think what we do for this depends on what we decide to do for specifying the subset of ops. Can we move forward with a boolean flag for now and change it later if we need to? We can change the name if you prefer, or maybe just pass an existing flag. But as far as I know, there's not a single existing flag that does what we want, because we don't want this to be triggered for internal mobile builds. @ljk53 any thoughts?

I lean towards having separate flags for now. We might still need decide how to specify ops whitelist, e.g. what if we want to choose between ops having same name? what if we want to use config file instead of encoded string? We can always make the "static dispatch" flag "intern_" flag and set it automatically later (IMO keeping it as separate intern flag is easier to understand anyway).

BTW, you probably only need modify .jenkins/pytorch/build.sh (see how DBUILD_CAFFE2_MOBILE is set there).

Is the plan for PyTorch Mobile to like, completely ignore OpenGL, Metal, Vulkan, etc?

@ajtulloch when we get there can we use switch-case or hashmap-of-function-pointers approach? The first step is to get rid of huge aten vtable which blocks linker from striping out unused code...

Static dispatch doesn't imply only CPU - the code below already generates the switch statements (which is good). We should add a configurable filter on dispatch ids (basically selecting which devices to compile), but it can be done separately.

CMakeLists.txt

aten/src/ATen/ScalarOps.h

aten/src/ATen/core/TensorMethods.h

smessmer · 2019-06-28T21:08:33Z

c10/core/TensorImpl.cpp

 #ifndef CAFFE2_FB_LIMITED_MOBILE_CAPABILITY

+#ifdef USE_STATIC_DISPATCH
+thread_local bool NonVariableTypeMode_enabled = true;


Why is variable type mode different depending on static dispatch?

For mobile, we're never going through Variable code, but we're passing Variables through TH methods that expect tensor. Specifically checked_tensor_unwrap does an is_variable() check, and it'll fail because NonVariableTypeMode is always off, even though we are never going through Variable code. Because our constraints are inference-only and never going through VariableType, I thought it would make sense to just set NonVariableTypeMode and keep it.

@yf225 Do you have any thoughts on this?

After @yf225's Variable/Tensor unification work, do we still need keep is_variable() check?

Earlier (before Will landed his work) when I tried removing virtual methods on types & variable, I figured it mostly worked fine without is_variable() check, only need keep a few variable overridden methods virtual:

ljk53@2f795a5
ljk53@cc8a178

It's not needed by mobile inference for now but we are still discussing whether we need variable/autodiff for federate on mobile in the future.

It's a bit scary - I don't say we will add autograd on mobile build, but coupling with static dispatch like that is suspicious. Can't we just have a lightweight guard (at::AutoNonVariableTypeMode) in every method call?

Probably even better would be to just make sure we don't produce Variables at all - e.g. stub out factory functions if autograd is not compiled in (do we have a dedicated flag for turning off autograd?)

Add static dispatch mode to reduce mobile code size gh-metadata: pytorch pytorch 22335 gh/li-roy/36/head

dzhulgakov · 2019-08-13T17:12:20Z

.circleci/verbatim-sources/linux-build-defaults.yml

        fi

+        # dispatch aten ops statically for mobile
+        if [[ ${BUILD_ENVIRONMENT} == *"android"* ]]; then


we should also filter for iOS here, @xta0 - what is the right way to do it?

I think we don't have ios CI yet.

li-roy · 2019-08-13T19:01:08Z

@pytorchbot retest this please

Add static dispatch mode to reduce mobile code size gh-metadata: pytorch pytorch 22335 gh/li-roy/36/head

smessmer · 2019-08-20T01:14:50Z

looks good, thanks

dzhulgakov · 2019-08-20T01:49:56Z

tools/jit/gen_jit_dispatch.py

        .device(${device})
        .pinned_memory(${pin_memory});
-auto result_ = torch::${name}(${args_with_tensor_options});
+#ifdef USE_STATIC_DISPATCH


Is it to avoid creating variables? It might be better to put NoGrad guard on. Afaiu, we want to get rid of at:: factory functions eventually and just create always variables (cc @gchanan)

Yeah it's to avoid creating variables. I don't think NoGrad works without additional changes, torch:: will always create a variable.

I plan to look into optional build for autograd/variable functionality on top of this PR and make changes if I find anything - so it's fine as long as the static dispatching part works.

ljk53 · 2019-08-20T19:11:33Z

aten/src/ATen/function_wrapper.py

 TENSOR_METHOD_DEFINITION = CodeTemplate("""\
 inline ${return_type} Tensor::${api_name}(${method_formals}) const {
+#ifdef USE_STATIC_DISPATCH
+    ${mobile_method_body}


nit: you might want to call this "static_dispatch_method_body" to be consistent with the macro name?

ljk53 · 2019-08-20T19:12:36Z

aten/src/ATen/function_wrapper.py

    ('BFloat16', 'BFloat16', 'BFloat16AccrealNotDefined', True),
 ]

+mobile_backends = ['CPU', 'QuantizedCPU', 'SparseCPU']


static_dispatch_backends?

ljk53 · 2019-08-20T19:16:18Z

tools/jit/gen_jit_dispatch.py

        .device(${device})
        .pinned_memory(${pin_memory});
-auto result_ = torch::${name}(${args_with_tensor_options});
+#ifdef USE_STATIC_DISPATCH


I plan to look into optional build for autograd/variable functionality on top of this PR and make changes if I find anything - so it's fine as long as the static dispatching part works.

facebook-github-bot · 2019-08-20T20:41:02Z

@li-roy merged this pull request in 6824c90.

Summary: Pull Request resolved: pytorch/pytorch#22335 Test Plan: Imported from OSS Differential Revision: D16048264 Pulled By: li-roy fbshipit-source-id: ad1e50951273962a51bac7c25c3d2e5a588a730e

ezyang · 2019-08-20T21:27:04Z

This broke the named tensor master-only build. @zou3519 https://circleci.com/gh/pytorch/pytorch/2519587?utm_campaign=vcs-integration-link&utm_medium=referral&utm_source=github-build-link/console

Add static dispatch mode to reduce mobile code size

9e08dd8

pytorchbot added caffe2 module: android Related to Android support module: build Build system issues module: internals Related to internal abstractions in c10 and ATen module: operators labels Jun 28, 2019

li-roy requested review from dzhulgakov, ljk53 and smessmer June 28, 2019 09:18

royboy added 2 commits June 28, 2019 03:46

Update on "Add static dispatch mode to reduce mobile code size"

b8f999c

Add static dispatch mode to reduce mobile code size gh-metadata: pytorch pytorch 22335 gh/li-roy/36/head

Update on "Add static dispatch mode to reduce mobile code size"

5f56c0e

Add static dispatch mode to reduce mobile code size gh-metadata: pytorch pytorch 22335 gh/li-roy/36/head

pytorchbot added the module: ci Related to continuous integration label Jun 28, 2019

Update on "Add static dispatch mode to reduce mobile code size"

a9b6c75

Add static dispatch mode to reduce mobile code size gh-metadata: pytorch pytorch 22335 gh/li-roy/36/head

smessmer requested changes Jun 28, 2019

View reviewed changes

Update on "Add static dispatch mode to reduce mobile code size"

caa6e1d

Add static dispatch mode to reduce mobile code size gh-metadata: pytorch pytorch 22335 gh/li-roy/36/head

pytorchbot added the oncall: jit Add this issue/PR to JIT oncall triage queue label Aug 13, 2019

Update on "Add static dispatch mode to reduce mobile code size"

96eaf5a

Add static dispatch mode to reduce mobile code size gh-metadata: pytorch pytorch 22335 gh/li-roy/36/head

dzhulgakov reviewed Aug 13, 2019

View reviewed changes

royboy added 2 commits August 13, 2019 15:08

Update on "Add static dispatch mode to reduce mobile code size"

99953ad

Add static dispatch mode to reduce mobile code size gh-metadata: pytorch pytorch 22335 gh/li-roy/36/head

Update on "Add static dispatch mode to reduce mobile code size"

2a7ed20

Add static dispatch mode to reduce mobile code size gh-metadata: pytorch pytorch 22335 gh/li-roy/36/head

smessmer approved these changes Aug 20, 2019

View reviewed changes

dzhulgakov approved these changes Aug 20, 2019

View reviewed changes

ljk53 approved these changes Aug 20, 2019

View reviewed changes

facebook-github-bot closed this in 6824c90 Aug 20, 2019

zou3519 deleted the gh/li-roy/36/head branch August 20, 2019 19:21

facebook-github-bot added the merged label Aug 20, 2019

li-roy restored the gh/li-roy/36/head branch August 20, 2019 23:53

li-roy deleted the gh/li-roy/36/head branch August 20, 2019 23:54

li-roy restored the gh/li-roy/36/head branch August 20, 2019 23:56

li-roy deleted the gh/li-roy/36/head branch August 21, 2019 00:27

mruberry added the Merged label Oct 28, 2020

Conversation

li-roy commented Jun 28, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

smessmer left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

li-roy Jun 28, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

li-roy commented Aug 13, 2019

Uh oh!

smessmer commented Aug 20, 2019

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot commented Aug 20, 2019

Uh oh!

ezyang commented Aug 20, 2019

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

li-roy commented Jun 28, 2019 •

edited

Loading

li-roy Jun 28, 2019 •

edited

Loading