Fix CUDA mask indexing, CUDA graph outputs and the runtime output-buffer check

.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
@@ -168,7 +169,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 +180,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 +206,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 +246,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(

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)
 

backends/apple/coreml/partition/coreml_partitioner.py

@@ -63,6 +63,65 @@
 )
 
 
+# Ops that coremltools can carry through a tensor above rank 5, by rewriting a
+# reshape -> transpose -> reshape chain into rank 4 steps.
+_HIGH_RANK_LAYOUT_TARGETS = frozenset(
+    [
+        exir_ops.edge.aten.view_copy.default,
+        exir_ops.edge.aten.permute_copy.default,
+        exir_ops.edge.aten.clone.default,
+        exir_ops.edge.dim_order_ops._clone_dim_order.default,
+    ]
+)
+
+
+def _produces_tensor_above_rank_5(node: torch.fx.Node) -> bool:
+    val = node.meta.get("val")
+    vals = list(val) if isinstance(val, (list, tuple)) else [val]
+    return any(isinstance(v, torch.Tensor) and v.dim() > 5 for v in vals)
+
+
+def _touches_unsupported_high_rank_tensor(
+    node: torch.fx.Node, is_skipped: Callable[[torch.fx.Node], bool]
+) -> bool:
+    """Core ML tensors have at most rank 5.
+
+    coremltools 9.0 raises "Core ML only supports tensors with rank <= 5" for
+    any tensor above rank 5 that it cannot rewrite away. It only rewrites a
+    static, unbranched layout chain, so such a tensor is fine only when every
+    node connected to it through tensors above rank 5 is a delegated view,
+    permute or clone, and every node producing such a tensor has a static shape
+    and a single user.
+    """
+    seen, todo = set(), [node]
+    while todo:
+        n = todo.pop()
+        if n in seen:
+            continue
+        seen.add(n)
+        high_rank_inputs = [
+            arg for arg in n.all_input_nodes if _produces_tensor_above_rank_5(arg)
+        ]
+        high_rank_output = _produces_tensor_above_rank_5(n)
+        if not high_rank_inputs and not high_rank_output:
+            continue
+        if (
+            n.op != "call_function"
+            or n.target not in _HIGH_RANK_LAYOUT_TARGETS
+            or is_skipped(n)
+        ):
+            return True
+        if high_rank_output and (
+            len(n.users) != 1
+            or any(isinstance(d, torch.SymInt) for d in n.meta["val"].shape)
+        ):
+            return True
+        todo += high_rank_inputs
+        if high_rank_output:
+            todo += list(n.users)
+    return False
+
+
 _ARG_MIN_MAX_TARGETS = (
     torch.ops.aten.argmax.default,
     torch.ops.aten.argmin.default,
@@ -83,6 +142,18 @@ def _is_arg_min_max_over_flattened_input(node: torch.fx.Node) -> bool:
     return dim is None
 
 
+def _reads_or_produces_bfloat16(node: torch.fx.Node) -> bool:
+    """coremltools 9.0 has no bfloat16 type.
+
+    It raises ``KeyError: torch.bfloat16`` on any op that reads or produces a
+    bfloat16 tensor, instead of reporting the op as unsupported.
+    """
+    val = node.meta.get("val")
+    vals = list(val) if isinstance(val, (list, tuple)) else [val]
+    vals += [arg.meta.get("val") for arg in node.all_input_nodes]
+    return any(isinstance(v, torch.Tensor) and v.dtype == torch.bfloat16 for v in vals)
+
+
 def _is_view_op(op: torch._ops.OpOverload) -> bool:
     schema = op._schema
     if len(schema.arguments) == 0:
@@ -152,13 +223,32 @@ def should_override_support(self, node) -> bool:
             )
             return True
 
+        if _touches_unsupported_high_rank_tensor(
+            node,
+            lambda n: self.should_skip_op_for_delegation(
+                getattr(n.target, "__name__", "").lower()
+            ),
+        ):
+            self.log_once(
+                f"{node.target} on a tensor above rank 5 is not supported by "
+                "CoreML.  Overriding op support."
+            )
+            return True
+
         if _is_arg_min_max_over_flattened_input(node):
             self.log_once(
                 "torch.ops.aten.{argmax, argmin}.default with dim=None is "
                 "not supported by CoreML.  Overriding op support."
             )
             return True
 
+        if _reads_or_produces_bfloat16(node):
+            self.log_once(
+                f"{node.target} on a bfloat16 tensor is not supported by CoreML.  "
+                "Overriding op support."
+            )
+            return True
+
         # TODO: enable this after bugs in ExecuTorch's partitioner are fixed
         # # If lower_full_graph=False, do not partition nodes with symbolic args because it can result in symbolic args
         # # in the placeholders due to partitioning, which CoreML does not support
@@ -362,21 +452,26 @@ def ops_to_not_decompose(
             torch.ops.aten.repeat_interleave.self_int,
             torch.ops.aten.repeat_interleave.self_Tensor,
         ]
+        rejected = set()
         for node in ep.graph.nodes:
             if node.op == "call_function" and isinstance(
                 node.target, torch._ops.OpOverload
             ):
                 try:
-                    if (
-                        op_support.is_node_supported(None, node)
-                        and node.target not in do_not_decompose_blocklist
-                        and not _is_view_op(node.target)
-                    ):
-                        do_not_decompose.append(node.target)
+                    supported = op_support.is_node_supported(None, node)
                 except Exception as e:
                     # CoreML's op_support.is_node_supported will sometimes throw
                     # for unsupported ops, rather than returning False
                     self.log_once(
                         f"Encountered exception when checking node support, treating node as unsupported: {e}"
                     )
-        return do_not_decompose, None
+                    supported = False
+                if not supported:
+                    rejected.add(node.target)
+                elif node.target not in do_not_decompose_blocklist and not _is_view_op(
+                    node.target
+                ):
+                    do_not_decompose.append(node.target)
+        # Without a per-node filter, a preserved op stays whole at every call, so
+        # one rejected call would be left with no kernel.
+        return [op for op in do_not_decompose if op not in rejected], None