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 0f1b335e5b83f402bda2dcdd9ecb411e0b67c651)
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
script_module_v4.ptl
script_module_v5.ptl
script_module_v6.ptl
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
test_autodiff.cpp [JIT] Propagate requires_grad to autodiff subgraphs (#71666) 2022-01-28 18:57:36 +00: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_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
test_backend_lib.cpp
test_class_import.cpp
test_class_parser.cpp
test_class_type.cpp
test_cleanup_passes.cpp
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
test_create_autodiff_subgraphs.cpp
test_cs_debug_info_serialization.cpp
test_custom_class.cpp
test_custom_class_registrations.cpp
test_custom_class_registrations.h
test_custom_operators.cpp
test_dce.cpp
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
test_gpu.cpp Nvfuser code bump 2_1_2022 (#72127) 2022-02-15 00:43:16 +00: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_graph_executor.cpp
test_graph_iterator.cpp
test_inliner.cpp
test_interface.cpp
test_interpreter.cpp
test_interpreter_async.pt
test_ir.cpp
test_irparser.cpp
test_jit_logging_levels.cpp
test_jit_type.cpp
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_interpreter_direct.cpp [Pytorch Edge] Refactor Compatibility Stuff into own directory (#71432) 2022-01-20 19:38:41 +00:00
test_lite_trainer.cpp
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
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
test_qualified_name.cpp
test_save_load.cpp [PyTorch] Adopt IValue::toTupleRef() where obvious (#65505) 2021-11-02 10:22:18 -07:00
test_schema_matching.cpp
test_script_profile.cpp
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
test_subgraph_matcher.cpp
test_subgraph_rewriter.cpp
test_subgraph_utils.cpp
test_union.cpp
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
test_utils.h
tests_setup.py
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*'