[vmap] Prepend "legacy" to files for old vmap implementation (#90324)

We have an older torch.vmap implementation. It is no longer supported.
It still needs to exist somewhere for the sake of BC with
torch.autograd.functional.

This PR makes it clear what files are meant for implementing the old
vmap implementation. I've seen a couple of PRs recently adding support
for the old vmap implementation, so this will lessen the confusion.

Test Plan:
- CI
Pull Request resolved: https://github.com/pytorch/pytorch/pull/90324
Approved by: https://github.com/samdow

.jenkins/pytorch/test.sh

@@ -237,7 +237,7 @@ test_dynamo_shard() {
       test_python_dispatch \
       test_fx \
       test_package \
-      test_vmap \
+      test_legacy_vmap \
     --shard "$1" "$NUM_TEST_SHARDS" \
     --verbose
   assert_git_not_dirty

aten/src/ATen/LegacyBatchedFallback.cpp

@@ -1,7 +1,7 @@
 #include <ATen/Context.h>
-#include <ATen/BatchedFallback.h>
+#include <ATen/LegacyBatchedFallback.h>
 #include <ATen/MatrixRef.h>
-#include <ATen/VmapTransforms.h>
+#include <ATen/LegacyVmapTransforms.h>
 #include <ATen/core/dispatch/Dispatcher.h>
 #include <c10/util/accumulate.h>
 #include <c10/util/llvmMathExtras.h>

aten/src/ATen/LegacyBatchedTensorImpl.cpp

@@ -1,4 +1,4 @@
-#include <ATen/BatchedTensorImpl.h>
+#include <ATen/LegacyBatchedTensorImpl.h>
 
 #include <ATen/WrapDimUtils.h>
 #include <c10/util/Exception.h>

aten/src/ATen/LegacyBatchingRegistrations.cpp

@@ -1,7 +1,7 @@
 #include <torch/library.h>
 #include <ATen/RedispatchFunctions.h>
-#include <ATen/VmapTransforms.h>
-#include <ATen/BatchedFallback.h>
+#include <ATen/LegacyVmapTransforms.h>
+#include <ATen/LegacyBatchedFallback.h>
 #include <ATen/native/ResizeCommon.h>
 #include <ATen/ATen.h>
 #include <ATen/core/IListRef.h>

aten/src/ATen/LegacyVmapMode.cpp

@@ -1,4 +1,4 @@
-#include <ATen/VmapMode.h>
+#include <ATen/LegacyVmapMode.h>
 
 namespace at {
 namespace impl {

aten/src/ATen/LegacyVmapTransforms.cpp

@@ -1,4 +1,4 @@
-#include <ATen/VmapTransforms.h>
+#include <ATen/LegacyVmapTransforms.h>
 #include <ATen/ATen.h>
 #include <ATen/core/IListRef.h>
 #include <c10/util/irange.h>

aten/src/ATen/LegacyVmapTransforms.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <ATen/BatchedTensorImpl.h>
+#include <ATen/LegacyBatchedTensorImpl.h>
 #include <ATen/core/IListRef.h>
 
 namespace at {

aten/src/ATen/native/LegacyBatching.cpp

@@ -1,7 +1,7 @@
 #include <ATen/core/Tensor.h>
-#include <ATen/BatchedTensorImpl.h>
+#include <ATen/LegacyBatchedTensorImpl.h>
 #include <ATen/WrapDimUtils.h>
-#include <ATen/VmapTransforms.h>
+#include <ATen/LegacyVmapTransforms.h>
 
 namespace at { namespace native {
 

aten/src/ATen/test/CMakeLists.txt

@@ -45,7 +45,7 @@ list(APPEND ATen_CPU_TEST_SRCS
   ${CMAKE_CURRENT_SOURCE_DIR}/undefined_tensor_test.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/variant_test.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/verify_api_visibility.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/vmap_test.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/legacy_vmap_test.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/weakref_test.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/wrapdim_test.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/xla_tensor_test.cpp

aten/src/ATen/test/legacy_vmap_test.cpp

@@ -1,8 +1,8 @@
 #include <gtest/gtest.h>
 
 #include <ATen/ATen.h>
-#include <ATen/BatchedTensorImpl.h>
-#include <ATen/VmapTransforms.h>
+#include <ATen/LegacyBatchedTensorImpl.h>
+#include <ATen/LegacyVmapTransforms.h>
 #include <c10/util/irange.h>
 
 using namespace at;

aten/tools/run_tests.sh

@@ -26,7 +26,7 @@ fi
 ./Dict_test
 ./NamedTensor_test
 ./cpu_generator_test
-./vmap_test
+./legacy_vmap_test
 ./operators_test
 if [[ -x ./cudnn_test ]]; then
   ./cudnn_test

build_variables.bzl

@@ -1028,7 +1028,7 @@ def glob_libtorch_python_sources(gencode_pattern = ":generate-code[{}]"):
 
 aten_cpu_source_non_codegen_list = [
     "aten/src/ATen/AccumulateType.cpp",
-    "aten/src/ATen/BatchedTensorImpl.cpp",
+    "aten/src/ATen/LegacyBatchedTensorImpl.cpp",
     "aten/src/ATen/CPUGeneratorImpl.cpp",
     "aten/src/ATen/Context.cpp",
     "aten/src/ATen/DLConvertor.cpp",
@@ -1062,8 +1062,8 @@ aten_cpu_source_non_codegen_list = [
     "aten/src/ATen/FuncTorchTLS.cpp",
     "aten/src/ATen/Utils.cpp",
     "aten/src/ATen/Version.cpp",
-    "aten/src/ATen/VmapMode.cpp",
-    "aten/src/ATen/VmapTransforms.cpp",
+    "aten/src/ATen/LegacyVmapMode.cpp",
+    "aten/src/ATen/LegacyVmapTransforms.cpp",
     "aten/src/ATen/core/BackendSelectFallbackKernel.cpp",
     "aten/src/ATen/core/DeprecatedTypeProperties.cpp",
     "aten/src/ATen/core/DeprecatedTypePropertiesRegistry.cpp",
@@ -1292,7 +1292,7 @@ aten_native_source_non_codegen_list = [
     "aten/src/ATen/native/AveragePool3d.cpp",
     "aten/src/ATen/native/BatchLinearAlgebra.cpp",
     "aten/src/ATen/native/BatchLinearAlgebraKernel.cpp",
-    "aten/src/ATen/native/Batching.cpp",
+    "aten/src/ATen/native/LegacyBatching.cpp",
     "aten/src/ATen/native/BinaryOps.cpp",
     "aten/src/ATen/native/Blas.cpp",
     "aten/src/ATen/native/BlasKernel.cpp",

torch/csrc/Module.cpp

@@ -9,10 +9,10 @@
 #include <ATen/ATen.h>
 #include <ATen/DLConvertor.h>
 #include <ATen/ExpandUtils.h>
+#include <ATen/LegacyVmapMode.h>
 #include <ATen/LinalgBackend.h>
 #include <ATen/Parallel.h>
 #include <ATen/Utils.h>
-#include <ATen/VmapMode.h>
 #include <ATen/core/Vitals.h>
 #include <ATen/dlpack.h>
 #include <ATen/native/ConvUtils.h>

torch/csrc/autograd/FunctionsManual.cpp

@@ -5,9 +5,9 @@
 
 #include <ATen/ATen.h>
 #include <ATen/AccumulateType.h>
-#include <ATen/BatchedTensorImpl.h>
 #include <ATen/Dispatch.h>
 #include <ATen/ExpandUtils.h>
+#include <ATen/LegacyBatchedTensorImpl.h>
 #include <ATen/ScalarOps.h>
 #include <ATen/SparseCsrTensorUtils.h>
 #include <ATen/SparseTensorUtils.h>

torch/csrc/autograd/functions/accumulate_grad.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <ATen/BatchedTensorImpl.h>
+#include <ATen/LegacyBatchedTensorImpl.h>
 #include <ATen/TensorOperators.h>
 #include <torch/csrc/Export.h>
 #include <torch/csrc/autograd/function.h>

torch/csrc/autograd/input_buffer.cpp

@@ -1,6 +1,6 @@
 #include <torch/csrc/autograd/input_buffer.h>
 
-#include <ATen/BatchedTensorImpl.h>
+#include <ATen/LegacyBatchedTensorImpl.h>
 #include <ATen/SparseCsrTensorUtils.h>
 #include <ATen/SparseTensorUtils.h>
 #include <ATen/TensorOperators.h>

torch/csrc/autograd/python_engine.cpp

@@ -1,7 +1,7 @@
 #include <torch/csrc/autograd/python_engine.h>
 
-#include <ATen/BatchedTensorImpl.h>
-#include <ATen/VmapMode.h>
+#include <ATen/LegacyBatchedTensorImpl.h>
+#include <ATen/LegacyVmapMode.h>
 #include <c10/util/irange.h>
 #include <pybind11/pybind11.h>
 #include <torch/csrc/DynamicTypes.h>