Index out of range in TransformReplay::replay()

[tlemo]

pytorch/torch/csrc/jit/codegen/cuda/transform_replay.cpp

Lines 113 to 116 in 172bcdb

	// Make sure both merge axes are real.
	TORCH_INTERNAL_ASSERT(
	axis_map[axis] != -1 && axis_map[axis + 1] != -1,
	"During transformation replay attempted to merge an imaginary axis.");

If I add a check there I see JitTest.GPU_FusionSimplePWise_CUDA failing:

[ RUN ] JitTest.GPU_FusionSimplePWise_CUDA
unknown file: error: C++ exception with description "axis + 1 < axis_map.size() INTERNAL ASSERT FAILED at ..\torch\csrc\jit\codegen\cuda\transform_replay.cpp:114, please report a bug to PyTorch. (replay at ..\torch\csrc\jit\codegen\cuda\transform_replay.cpp:114)
00007FF95523E77200007FF95523E730 c10.dll!std::_Invoker_functor::_Call<<lambda_caff7cead0c0e49cc41a504cb708e59a> &> [C:\Program Files (x86)\Microsoft Visual Studio\2019\Community\VC\Tools\MSVC\14.24.28314\include\type_traits @ 1579]
00007FF95523EBC200007FF95523EB80 c10.dll!std::invoke<<lambda_caff7cead0c0e49cc41a504cb708e59a> &> [C:\Program Files (x86)\Microsoft Visual Studio\2019\Community\VC\Tools\MSVC\14.24.28314\include\type_traits @ 1579]
00007FF95523E7D200007FF95523E790 c10.dll!std::_Invoker_ret<std::basic_string<char,std::char_traits,std::allocator >,0>::_Call<<lambda_caff7cead0c0e49cc41a504cb708e59a> &> [C:\Program Files (x86)\Microsoft Visual Studio\2019\Community\VC\Tools\MSVC\14.24.
00007FF95523F47100007FF95523F430 c10.dll!std::_Func_impl_no_alloc<<lambda_caff7cead0c0e49cc41a504cb708e59a>,std::basic_string<char,std::char_traits,std::allocator > >::_Do_call [C:\Program Files (x86)\Microsoft Visual Studio\2019\Community\VC\Tools\MSVC
00007FF95523F24500007FF95523F1E0 c10.dll!std::_Func_class<std::basic_string<char,std::char_traits,std::allocator > >::operator() [C:\Program Files (x86)\Microsoft Visual Studio\2019\Community\VC\Tools\MSVC\14.24.28314\include\functional @ 969]
00007FF95523D49200007FF95523D440 c10.dll!c10::Error::Error [C:\Users\lemo\work\external\pytorch\pytorch\c10\util\Logging.cpp @ 67]
00007FF909AAC9CC00007FF909AAC280 torch_cuda.dll!torch::jit::fuser::TransformReplay::replay [C:\Users\lemo\work\external\pytorch\pytorch\torch\csrc\jit\codegen\cuda\transform_replay.cpp @ 114]
00007FF909AAA6FA00007FF909AAA4E0 torch_cuda.dll!torch::jit::fuser::TransformIter::replay [C:\Users\lemo\work\external\pytorch\pytorch\torch\csrc\jit\codegen\cuda\transform_iter.cpp @ 105]
00007FF909AAA94C00007FF909AAA890 torch_cuda.dll!torch::jit::fuser::TransformIter::runReplay [C:\Users\lemo\work\external\pytorch\pytorch\torch\csrc\jit\codegen\cuda\transform_iter.cpp @ 115]
00007FF909AADA0000007FF909AAD4C0 torch_cuda.dll!torch::jit::fuser::TransformReplay::runReplay [C:\Users\lemo\work\external\pytorch\pytorch\torch\csrc\jit\codegen\cuda\transform_replay.cpp @ 308]
00007FF909AADC7000007FF909AADC10 torch_cuda.dll!torch::jit::fuser::TransformReplay::runReplay [C:\Users\lemo\work\external\pytorch\pytorch\torch\csrc\jit\codegen\cuda\transform_replay.cpp @ 323]
00007FF909AADCFB00007FF909AADCA0 torch_cuda.dll!torch::jit::fuser::TransformReplay::replay [C:\Users\lemo\work\external\pytorch\pytorch\torch\csrc\jit\codegen\cuda\transform_replay.cpp @ 335]
00007FF909AA765700007FF909AA6E80 torch_cuda.dll!torch::jit::fuser::TensorView::computeAt [C:\Users\lemo\work\external\pytorch\pytorch\torch\csrc\jit\codegen\cuda\tensor_view.cpp @ 178]
00007FF909AA73C900007FF909AA6E80 torch_cuda.dll!torch::jit::fuser::TensorView::computeAt [C:\Users\lemo\work\external\pytorch\pytorch\torch\csrc\jit\codegen\cuda\tensor_view.cpp @ 139]
00007FF7339436A100007FF733943090 test_jit.exe!torch::jit::testGPU_FusionSimplePWise [C:\Users\lemo\work\external\pytorch\pytorch\test\cpp\jit\test_gpu.cpp @ 979]
00007FF7337F89D300007FF7337F89B0 test_jit.exe!torch::jit::JitTest_GPU_FusionSimplePWise_CUDA_Test::TestBody [C:\Users\lemo\work\external\pytorch\pytorch\test\cpp\jit\gtest.cpp @ 19]
00007FF733ADE94100007FF733ADE910 test_jit.exe!testing::internal::HandleSehExceptionsInMethodIfSupportedtesting::Test,void [C:\Users\lemo\work\external\pytorch\pytorch\third_party\googletest\googletest\src\gtest.cc @ 2428]
00007FF733ADE64100007FF733ADE5D0 test_jit.exe!testing::internal::HandleExceptionsInMethodIfSupportedtesting::Test,void [C:\Users\lemo\work\external\pytorch\pytorch\third_party\googletest\googletest\src\gtest.cc @ 2479]
00007FF733AC2F8300007FF733AC2ED0 test_jit.exe!testing::test::Run [C:\Users\lemo\work\external\pytorch\pytorch\third_party\googletest\googletest\src\gtest.cc @ 2524]
00007FF733AC3B1E00007FF733AC3A40 test_jit.exe!testing::TestInfo::Run [C:\Users\lemo\work\external\pytorch\pytorch\third_party\googletest\googletest\src\gtest.cc @ 2697]
00007FF733AC42A600007FF733AC41B0 test_jit.exe!testing::TestCase::Run [C:\Users\lemo\work\external\pytorch\pytorch\third_party\googletest\googletest\src\gtest.cc @ 2812]
00007FF733ACB51E00007FF733ACB1C0 test_jit.exe!testing::internal::UnitTestImpl::RunAllTests [C:\Users\lemo\work\external\pytorch\pytorch\third_party\googletest\googletest\src\gtest.cc @ 5178]
00007FF733ADEA8100007FF733ADEA50 test_jit.exe!testing::internal::HandleSehExceptionsInMethodIfSupportedtesting::internal::UnitTestImpl,bool [C:\Users\lemo\work\external\pytorch\pytorch\third_party\googletest\googletest\src\gtest.cc @ 2428]
00007FF733ADE8A100007FF733ADE830 test_jit.exe!testing::internal::HandleExceptionsInMethodIfSupportedtesting::internal::UnitTestImpl,bool [C:\Users\lemo\work\external\pytorch\pytorch\third_party\googletest\googletest\src\gtest.cc @ 2479]
00007FF733AC499300007FF733AC4870 test_jit.exe!testing::UnitTest::Run [C:\Users\lemo\work\external\pytorch\pytorch\third_party\googletest\googletest\src\gtest.cc @ 4786]

[csarofeen]

All these operations are now in transform_iter.cpp. They were moved to support reductions/rfactor. I looked through all the replay merge functions and no longer see an assert like this. I think we should be good, but would recommend, if you'd like, to review our changes there.