CUDA backend: per-program weights file names and antialiased resize

.ci/scripts/wheel/test_shared_libraries.py

@@ -104,6 +104,8 @@
     "aoti_torch_cuda_randint_low_out",
     "aoti_torch_cuda_sort_stable",
     "aoti_torch_cuda_guard_set_index",
+    "aoti_torch_cuda_index_Tensor",
+    "aoti_torch_index_put_out",
 )
 
 _QUANTIZED_KERNEL_SYMBOLS = (

backends/aoti/aoti_backend.py

@@ -7,6 +7,7 @@
 import contextlib
 import hashlib
 import os
+import re
 import typing
 from abc import ABC, abstractmethod
 from enum import Enum
@@ -27,6 +28,29 @@
 
 class COMPILE_SPEC_KEYS(Enum):
     METHOD_NAME = "method_name"
+    DATA_NAME = "data_name"
+
+
+# The name becomes a file name on every platform, so keep it short, to portable
+# characters, and clear of Windows device names such as NUL. Without re.ASCII,
+# IGNORECASE lets [A-Za-z] match letters like the Kelvin sign. A leading '-'
+# reads as an option on a command line. Some writers add .ptd only when it is
+# missing, so "encoder" and "encoder.ptd" could both write encoder.ptd.
+_DATA_NAME_PATTERN = re.compile(
+    r"(?!(?:CON|PRN|AUX|NUL|COM\d|LPT\d)(?:\.|$))(?!.*\.ptd$)"
+    r"[A-Za-z0-9_][A-Za-z0-9_.-]{0,199}",
+    re.IGNORECASE | re.ASCII,
+)
+
+
+def _checked_data_name(data_name: str) -> str:
+    if not _DATA_NAME_PATTERN.fullmatch(data_name):
+        raise ValueError(
+            f"Invalid data name {data_name!r}: use up to 200 ASCII letters, "
+            "digits, '_', '-' and '.', do not start with '.' or '-', do not end "
+            "with '.ptd', and avoid Windows device names such as NUL"
+        )
+    return data_name
 
 
 @experimental(
@@ -168,7 +192,8 @@ def release_moved_tensors(
     def collect_unsupported_fallback_kernels(cls, missing_fallback_kernels: Set[str]):
         """
         Context manager to collect unsupported fallback kernels during compilation.
-        Monitors both extern kernel calls and runtime lookup.
+        Monitors extern kernel calls, runtime lookup, and the scatter and
+        index_put fallbacks, which call their kernels directly.
         """
         supported_kernels = cls.get_supported_fallback_kernels()
 
@@ -178,6 +203,10 @@ def collect_unsupported_fallback_kernels(cls, missing_fallback_kernels: Set[str]
         original_generate_fallback_kernel_with_runtime_lookup_aot = (
             CppWrapperCpu.generate_fallback_kernel_with_runtime_lookup_aot
         )
+        original_generate_scatter_fallback = CppWrapperCpu._generate_scatter_fallback
+        original_generate_index_put_fallback = (
+            CppWrapperCpu._generate_index_put_fallback
+        )
 
         def generate_c_shim_extern_kernel_call_and_collect_unsupported_kernels(
             self, kernel: str, *args: Any, **kwargs: Any
@@ -200,10 +229,36 @@ def generate_fallback_kernel_with_runtime_lookup_aot_and_collect_unsupported_ker
                 self, op_overload, *args, **kwargs
             )
 
+        def collect_kernels_it_writes(original: Any) -> Any:
+            # These fallbacks write their shim call directly, so read the
+            # kernel names from the lines they write.
+            def generate_and_collect_unsupported_kernels(
+                self, *args: Any, **kwargs: Any
+            ) -> None:
+                lines: List[str] = []
+                writeline = self.writeline
+                self.writeline = lambda line: (lines.append(str(line)), writeline(line))
+                try:
+                    return original(self, *args, **kwargs)
+                finally:
+                    self.writeline = writeline
+                    for line in lines:
+                        for kernel in re.findall(r"\b(aoti_torch_\w+)\(", line):
+                            if kernel not in supported_kernels:
+                                missing_fallback_kernels.add(kernel)
+
+            return generate_and_collect_unsupported_kernels
+
         CppWrapperCpu.generate_c_shim_extern_kernel_call = (
             generate_c_shim_extern_kernel_call_and_collect_unsupported_kernels
         )
         CppWrapperCpu.generate_fallback_kernel_with_runtime_lookup_aot = generate_fallback_kernel_with_runtime_lookup_aot_and_collect_unsupported_kernels
+        CppWrapperCpu._generate_scatter_fallback = collect_kernels_it_writes(
+            original_generate_scatter_fallback
+        )
+        CppWrapperCpu._generate_index_put_fallback = collect_kernels_it_writes(
+            original_generate_index_put_fallback
+        )
 
         try:
             yield
@@ -214,6 +269,12 @@ def generate_fallback_kernel_with_runtime_lookup_aot_and_collect_unsupported_ker
             CppWrapperCpu.generate_fallback_kernel_with_runtime_lookup_aot = (
                 original_generate_fallback_kernel_with_runtime_lookup_aot
             )
+            CppWrapperCpu._generate_scatter_fallback = (
+                original_generate_scatter_fallback
+            )
+            CppWrapperCpu._generate_index_put_fallback = (
+                original_generate_index_put_fallback
+            )
 
     @classmethod
     def preprocess(
@@ -225,6 +286,9 @@ def preprocess(
         Preprocess the edge program and compile it using AOTInductor.
         Weights are always separated from the SO file.
         """
+        # First, so that a backend keeping its data in the .pte refuses a data
+        # name before anything is compiled.
+        data_name = cls.data_name_from_compile_specs(compile_specs)
         device_name = cls.get_device_name()
         decomposition_table = cls.get_decomposition_table()
         options = cls.get_aoti_compile_options(compile_specs)
@@ -329,9 +393,7 @@ def preprocess(
         named_data_store.add_named_data(so_blob_key, so_data, 1, None)
         # Determine whether to save named data externally based on backend setting
         # External: save to separate .ptd file, otherwise merge with .pte file
-        external_tag = (
-            f"aoti_{device_name}_blob" if cls.save_data_externally() else None
-        )
+        external_tag = data_name if cls.save_data_externally() else None
 
         named_data_store.add_named_data(weights_blob_key, blob_data, 1, external_tag)
 
@@ -366,6 +428,38 @@ def generate_method_name_compile_spec(
             method_name.encode("utf-8"),
         )
 
+    @classmethod
+    def generate_data_name_compile_spec(cls, data_name: str) -> CompileSpec:
+        """
+        Generate a CompileSpec that names the external data file, written as
+        ``<data_name>.ptd``, so several programs can share a folder.
+        """
+        cls._check_saves_data_externally()
+        return CompileSpec(
+            COMPILE_SPEC_KEYS.DATA_NAME.value,
+            _checked_data_name(data_name).encode("utf-8"),
+        )
+
+    @classmethod
+    def data_name_from_compile_specs(cls, compile_specs: List[CompileSpec]) -> str:
+        """
+        The external data file name, without ``.ptd``, set by
+        generate_data_name_compile_spec, or ``aoti_<device>_blob`` by default.
+        """
+        for spec in compile_specs:
+            if spec.key == COMPILE_SPEC_KEYS.DATA_NAME.value:
+                cls._check_saves_data_externally()
+                return _checked_data_name(spec.value.decode("utf-8"))
+        return f"aoti_{cls.get_device_name()}_blob"
+
+    @classmethod
+    def _check_saves_data_externally(cls) -> None:
+        if not cls.save_data_externally():
+            raise ValueError(
+                f"{cls.__name__} keeps its data in the .pte file, so it takes no "
+                "data name"
+            )
+
     @classmethod
     def method_name_from_compile_specs(
         cls,

backends/aoti/common_shims_slim.cpp

@@ -36,6 +36,19 @@ AOTITorchError aoti_torch_get_sizes(Tensor* tensor, int64_t** ret_sizes) {
   return Error::Ok;
 }
 
+AOTITorchError
+aoti_torch_get_size(Tensor* tensor, int64_t d, int64_t* ret_size) {
+  if (tensor == nullptr || ret_size == nullptr) {
+    return Error::InvalidArgument;
+  }
+  const auto ndim = static_cast<int64_t>(tensor->dim());
+  if (d < -ndim || d >= ndim) {
+    return Error::InvalidArgument;
+  }
+  *ret_size = tensor->size(d);
+  return Error::Ok;
+}
+
 AOTITorchError aoti_torch_get_strides(Tensor* tensor, int64_t** ret_strides) {
   if (tensor == nullptr || ret_strides == nullptr) {
     return Error::InvalidArgument;

backends/aoti/common_shims_slim.h

@@ -39,6 +39,11 @@ aoti_torch_get_data_ptr(Tensor* tensor, void** ret_data_ptr);
 AOTI_SHIM_EXPORT AOTITorchError
 aoti_torch_get_sizes(Tensor* tensor, int64_t** ret_sizes);
 
+// Generated code reads the size of a result whose shape depends on the data
+// through this, one dim at a time.
+AOTI_SHIM_EXPORT AOTITorchError
+aoti_torch_get_size(Tensor* tensor, int64_t d, int64_t* ret_size);
+
 AOTI_SHIM_EXPORT AOTITorchError
 aoti_torch_get_strides(Tensor* tensor, int64_t** ret_strides);
 

backends/aoti/passes/replace_view_copy_with_view.py

@@ -31,6 +31,12 @@
     # natively by inductor, so map it here like the other view-copy ops.
     torch.ops.aten.split_copy.Tensor: torch.ops.aten.split.Tensor,
     ops.edge.aten.split_copy.Tensor: ops.edge.aten.split.Tensor,
+    # Edge turns ``detach`` into ``detach_copy``, which has no c-shim either.
+    # Without autograd it is ``alias_copy``, which Inductor can decompose,
+    # unlike ``detach_copy``. A view would let two outputs share one handle,
+    # and ``clone`` fails Edge verification in dim order mode.
+    torch.ops.aten.detach_copy.default: torch.ops.aten.alias_copy.default,
+    ops.edge.aten.detach_copy.default: ops.edge.aten.alias_copy.default,
 }
 
 
@@ -40,19 +46,26 @@ class ReplaceViewCopyWithViewPass(ExportPass):
     def call(self, graph_module: fx.GraphModule) -> PassResult:
         graph_changed = False
 
-        for node in graph_module.graph.nodes:
-            if node.op != "call_function" or node.target not in _VIEW_TARGETS:
+        # Control flow branches, such as those of torch.cond, are submodules.
+        for module in graph_module.modules():
+            if not isinstance(module, fx.GraphModule):
                 continue
 
-            if self._has_blocking_user(node, node.users.keys()):
-                continue
+            module_changed = False
+            for node in module.graph.nodes:
+                if node.op != "call_function" or node.target not in _VIEW_TARGETS:
+                    continue
+
+                if self._has_blocking_user(node, node.users.keys()):
+                    continue
 
-            node.target = _VIEW_TARGETS[node.target]
-            graph_changed = True
+                node.target = _VIEW_TARGETS[node.target]
+                module_changed = True
 
-        if graph_changed:
-            graph_module.graph.lint()
-            graph_module.recompile()
+            if module_changed:
+                module.graph.lint()
+                module.recompile()
+                graph_changed = True
 
         return PassResult(graph_module, graph_changed)