pytorch/test/cpp/jit
Mike Iovine d1c5f9e439 [JIT][SR] Introduce prim::IfThenElse (#72587)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/72587

This pattern frequently appears in a few graphs:

```
%result = prim::If(%condition)
  block0():
    -> (%a)
  block1():
    -> (%b)
```

This is slow, particularly in static runtime. Static runtime creates memory planners/block runners for each sub-block, which eats up a lot of memory and introduces a lot of extra overhead for this relatively simple operation.

This diff introduces a new op that replaces nodes like the above with a single op meant to act like a ternary operator:

```
%result = prim::IfThenElse(%condition, %a, %b)
```

Test Plan: New unit tests

Reviewed By: eellison

Differential Revision: D34091789

fbshipit-source-id: eb6a8c460c39b4c019a1f4ab1f3f1e5b6edc400c
(cherry picked from commit 0f1b335e5b)
2022-02-17 18:22:48 +00:00
..
upgrader_models [Operator Versioning][Edge] Add old models and unittest (#67726) 2021-12-01 18:46:30 -08:00
__init__.py
CMakeLists.txt [JIT][SR] Introduce prim::IfThenElse (#72587) 2022-02-17 18:22:48 +00:00
README.md port all JIT tests to gtest (#45264) 2020-09-25 11:37:43 -07:00
script_module_v4.ptl [PyTorch][Edge] Add binarires for unittests (#57039) 2021-04-27 20:46:34 -07:00
script_module_v5.ptl [PyTorch][Edge] Add binarires for unittests (#57039) 2021-04-27 20:46:34 -07:00
script_module_v6.ptl [Pytorch Edge] Model Ops compatibility api (#57501) 2021-05-24 12:00:06 -07:00
test_add_if_then_else.cpp [JIT][SR] Introduce prim::IfThenElse (#72587) 2022-02-17 18:22:48 +00:00
test_alias_analysis.cpp Codegen: python_torch_functions only include relevant operators (#68693) 2022-01-21 15:37:06 +00:00
test_argument_spec.cpp [jit] Remove graph() call from abstract Function interface. (#65967) 2021-10-27 11:54:26 -07:00
test_autodiff.cpp [JIT] Propagate requires_grad to autodiff subgraphs (#71666) 2022-01-28 18:57:36 +00:00
test_backend_compiler_lib.cpp [Pytorch Edge] Minor improve documentation in test_backend_with_compiler 2022-02-14 23:44:46 +00:00
test_backend_compiler_preprocess.cpp [jit] Remove graph() call from abstract Function interface. (#65967) 2021-10-27 11:54:26 -07:00
test_backend_lib.cpp Disable avoid-non-const-global-variables lint check (#62008) 2021-07-22 18:04:40 -07:00
test_backend.cpp (1/2) Make TorchScript Preserve Fully Qualified Class Name for Python Exceptions: backend change 2022-01-27 03:29:40 +00:00
test_class_import.cpp Add SourceView which doesn't own source text as base class of Source (#65309) 2021-10-18 23:17:22 -07:00
test_class_parser.cpp Add SourceView which doesn't own source text as base class of Source (#65309) 2021-10-18 23:17:22 -07:00
test_class_type.cpp Disable avoid-non-const-global-variables lint check (#62008) 2021-07-22 18:04:40 -07:00
test_cleanup_passes.cpp Disable avoid-non-const-global-variables lint check (#62008) 2021-07-22 18:04:40 -07:00
test_code_template.cpp [AutoAccept][Codemod][FBSourceClangFormatLinter] Daily arc lint --take CLANGFORMAT 2022-01-14 06:10:57 -08:00
test_concat_opt.cpp [JIT] Combine concat nodes where possible (#67000) 2021-11-15 12:02:45 -08:00
test_constant_pooling.cpp Disable avoid-non-const-global-variables lint check (#62008) 2021-07-22 18:04:40 -07:00
test_create_autodiff_subgraphs.cpp Disable avoid-non-const-global-variables lint check (#62008) 2021-07-22 18:04:40 -07:00
test_cs_debug_info_serialization.cpp [Pytorch Delegated Backend] Save function name in debug info (#57481) 2021-05-25 13:19:02 -07:00
test_custom_class_registrations.cpp [clang-tidy] Exclude cppcoreguidelines-avoid-magic-numbers (#57841) 2021-05-07 20:02:33 -07:00
test_custom_class_registrations.h
test_custom_class.cpp Disable avoid-non-const-global-variables lint check (#62008) 2021-07-22 18:04:40 -07:00
test_custom_operators.cpp [jit] Reduce refcounting of Types (#65345) 2021-10-08 09:03:04 -07:00
test_dce.cpp Disable avoid-non-const-global-variables lint check (#62008) 2021-07-22 18:04:40 -07:00
test_exception.cpp (2/2) Make TorchScript Preserve Fully Qualified Class Name for Python Exceptions: frontend change (#72899) 2022-02-16 03:45:15 +00:00
test_flatbuffer.cpp Fixing few bugs in torch flatbuffer (#72349) 2022-02-05 00:25:27 +00:00
test_fuser.cpp Disable avoid-non-const-global-variables lint check (#62008) 2021-07-22 18:04:40 -07:00
test_gpu_shift.cpp Nvfuser code bump 2_1_2022 (#72127) 2022-02-15 00:43:16 +00:00
test_gpu_validator.h Nvfuser code bump 2_1_2022 (#72127) 2022-02-15 00:43:16 +00:00
test_gpu.cpp Nvfuser code bump 2_1_2022 (#72127) 2022-02-15 00:43:16 +00:00
test_graph_executor.cpp Disable avoid-non-const-global-variables lint check (#62008) 2021-07-22 18:04:40 -07:00
test_graph_iterator.cpp [jit] Rewrote DFS graph iterator to remove unnecessary local state (#61326) (#61980) 2021-07-27 11:50:20 -07:00
test_inliner.cpp [jit] Remove graph() call from abstract Function interface. (#65967) 2021-10-27 11:54:26 -07:00
test_interface.cpp Add SourceView which doesn't own source text as base class of Source (#65309) 2021-10-18 23:17:22 -07:00
test_interpreter_async.pt
test_interpreter.cpp [PyTorch Edge] Backport function for defaults args with out args, flag on (#63651) 2021-09-20 22:50:30 -07:00
test_ir.cpp Disable avoid-non-const-global-variables lint check (#62008) 2021-07-22 18:04:40 -07:00
test_irparser.cpp [jit] Reduce refcounting of Types (#65345) 2021-10-08 09:03:04 -07:00
test_jit_logging_levels.cpp Added jit log stream changing function and some refactor (#65768) 2021-09-30 23:25:11 -07:00
test_jit_type.cpp [jit] Reduce refcounting of Types (#65345) 2021-10-08 09:03:04 -07:00
test_lite_interpreter_direct.cpp [Pytorch Edge] Refactor Compatibility Stuff into own directory (#71432) 2022-01-20 19:38:41 +00:00
test_lite_interpreter.cpp [PyTorchEdge] backport v8 to v7 to support promoted ops as instruction (#71662) 2022-02-15 03:47:39 +00:00
test_lite_trainer.cpp Disable avoid-non-const-global-variables lint check (#62008) 2021-07-22 18:04:40 -07:00
test_load_upgraders.cpp Bump version number to 7 and compile old operators with old schema (#68358) 2022-01-05 23:57:22 -08:00
test_memory_dag.cpp Disable avoid-non-const-global-variables lint check (#62008) 2021-07-22 18:04:40 -07:00
test_misc.cpp [JIT] enable profiling optional tensors (#70532) 2022-02-08 22:52:26 +00:00
test_mobile_type_parser.cpp [jit][edge] Migrate base types to DynamicType on mobile. (#70233) 2022-01-11 13:53:29 -08:00
test_module_api.cpp Fix bug where frozen mod not used for OFI #68903 (#71436) 2022-01-27 23:37:50 +00:00
test_op_replacement.cpp Add graph op replacement pass (#69915) 2021-12-25 13:03:19 -08:00
test_peephole_optimize.cpp Disable avoid-non-const-global-variables lint check (#62008) 2021-07-22 18:04:40 -07:00
test_qualified_name.cpp Disable avoid-non-const-global-variables lint check (#62008) 2021-07-22 18:04:40 -07:00
test_save_load.cpp [PyTorch] Adopt IValue::toTupleRef() where obvious (#65505) 2021-11-02 10:22:18 -07:00
test_schema_matching.cpp [jit] Make operation call accept Stack& instead Stack* (#63414) 2021-08-30 11:49:20 -07:00
test_script_profile.cpp Reapply "[jit] Implement ScriptProfile to collect instruction profiles." (#58783) 2021-05-24 18:23:21 -07:00
test_shape_analysis.cpp Add support for permutting dynamic fusion group outputs to channels last format (#70656) 2022-01-12 09:11:34 -08:00
test_stack_opt.cpp [JIT] Add variadic stack op (#63578) 2021-08-24 08:20:54 -07:00
test_subgraph_matcher.cpp Disable avoid-non-const-global-variables lint check (#62008) 2021-07-22 18:04:40 -07:00
test_subgraph_rewriter.cpp keep output type after calling SubgraphRewriter (#65453) 2021-09-24 11:07:40 -07:00
test_subgraph_utils.cpp Disable avoid-non-const-global-variables lint check (#62008) 2021-07-22 18:04:40 -07:00
test_union.cpp [jit] Reduce refcounting of Types (#65345) 2021-10-08 09:03:04 -07:00
test_upgrader_utils.cpp Add option to load historic operators in IR when the operator is deprecated (#71148) 2022-01-12 11:07:04 -08:00
test_utils.cpp [jit] Make operation call accept Stack& instead Stack* (#63414) 2021-08-30 11:49:20 -07:00
test_utils.h Enhancement for smart serialization for out schemas (#63096) 2021-08-28 11:46:27 -07:00
tests_setup.py Revert D30279364: [codemod][lint][fbcode/c*] Enable BLACK by default 2021-08-12 11:45:01 -07:00
torch_python_test.cpp Remove WindowsTorchApiMacro.h in favor of Export.h (#69585) 2021-12-09 17:30:09 -08:00

JIT C++ Tests

Adding a new test

First, create a new test file. Test files should have be placed in this directory, with a name that starts with test_, like test_foo.cpp.

In general a single test suite

Add your test file to the JIT_TEST_SRCS list in test/cpp/jit/CMakeLists.txt.

A test file may look like:

#include <gtest/gtest.h>

using namespace ::torch::jit

TEST(FooTest, BarBaz) {
   // ...
}

// Append '_CUDA' to the test case name will automatically filter it out if CUDA
// is not compiled.
TEST(FooTest, NeedsAGpu_CUDA) {
   // ...
}

// Similarly, if only one GPU is detected, tests with `_MultiCUDA` at the end
// will not be run.
TEST(FooTest, NeedsMultipleGpus_MultiCUDA) {
   // ...
}

Building and running the tests

The following commands assume you are in PyTorch root.

# ... Build PyTorch from source, e.g.
python setup.py develop
# (re)build just the binary
ninja -C build bin/test_jit
# run tests
build/bin/test_jit --gtest_filter='glob_style_filter*'