OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2026-01-15 12:15:51 +00:00
Code Issues Actions 52 Packages Projects Releases Wiki Activity

Add dynamic shape support to AOT driver & compiler (#72995)

Browse Source 
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/72995

Add ability to specify input dimensions that need to be dynamic.
Example: if dim 115 can be dynamic in input sizes "1,115;1", then specify dynamic_dims as "115"

Also recompile and update CI models and some asm code as the old ones don't compile with compiler changes in context.cpp

Test Plan: - Compiles and runs BI Bytedoc model with and without dynamic inputs.

Reviewed By: ZolotukhinM

Differential Revision: D34233121

fbshipit-source-id: 35095e549ebd6d3bec98b9abb3f0764366a0ff6f
(cherry picked from commit 33166a9f9ac9194b5df0a35280b57708df255ebd)
This commit is contained in:
Priya Ramani
2022-02-23 19:56:22 -08:00
committed by PyTorch MergeBot
parent 5a7778c9a6
commit ac97e953b4
4 changed files with 104 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
binaries/aot_model_compiler.cc
View File

@@ -36,6 +36,10 @@ C10_DEFINE_string(
"Input memory format."
"If multiple inputs needed, use semicolon to separate."
"Supported values: contiguous, channels_last");
C10_DEFINE_string(
dynamic_dims,
"",
"Comma separated dimensions of input tensors that can be dynamic");
C10_DEFINE_string(method_name, "forward", "The name of the method.");
C10_DEFINE_string(
output_llvm,
@@ -68,6 +72,7 @@ c10::Dict<c10::IValue, c10::IValue> createCompileSpec() {
method_spec.insert("sizes", FLAGS_input_dims);
method_spec.insert("types", FLAGS_input_types);
method_spec.insert("memory_formats", FLAGS_input_memory_formats);
method_spec.insert("dynamic_sizes", FLAGS_dynamic_dims);
method_spec.insert("asmfile", FLAGS_output_llvm);
method_spec.insert("model_name", FLAGS_model_name);
method_spec.insert("model_version", FLAGS_model_version);