Inductor respects exact strides on custom ops by default

zou3519 · zou3519 · commit c3360aafbe0e · 2025-02-27T09:27:35.000-08:00
If a tag is not specified on a custom operator, then inductor will
assume that it needs exact strides.

Test Plan:
- tests + CI

[ghstack-poisoned]
diff --git a/test/inductor/test_triton_kernels.py b/test/inductor/test_triton_kernels.py
@@ -3457,7 +3457,6 @@ def impl2(x):
 
             lib.define(
                 "add_op(Tensor x, Tensor y) -> Tensor",
-                tags=[torch._C.Tag.needs_exact_strides],
             )
 
             def impl(x, y):
@@ -3471,7 +3470,6 @@ def meta(x, y):
 
             lib.define(
                 "add_out_op(Tensor x, Tensor y, Tensor(a!) out) -> ()",
-                tags=[torch._C.Tag.needs_exact_strides],
             )
 
             def impl_out(x, y, out):
diff --git a/torch/_inductor/graph.py b/torch/_inductor/graph.py
@@ -1124,21 +1124,19 @@ def call_function(self, target: Callable, args: Any, kwargs: dict[str, Any]) ->
                     error.operator_str(target, args, kwargs),
                 )
 
-                decided_constraint = require_contiguous
-
-                # use contiguous unless the (custom) op asks something else
-                # explicitly
-                if torch._C.Tag.needs_exact_strides in target.tags:
-                    decided_constraint = constrain_to_fake_tensors  # type: ignore[assignment]
+                if torch._library.utils.needs_exact_strides(target):
+                    decided_constraint = constrain_to_fake_tensors
                 elif torch._C.Tag.needs_fixed_stride_order in target.tags:
                     decided_constraint = constrain_to_fx_strides  # type: ignore[assignment]
                 elif torch._C.Tag.flexible_layout in target.tags:
                     decided_constraint = None  # type: ignore[assignment]
+                else:
+                    # for implicitly fallback ops, we conservatively requires
+                    # contiguous input since some eager kernels does not
+                    # support non-contiguous inputs. They may silently cause
+                    # accuracy problems. Check https://github.com/pytorch/pytorch/issues/140452
+                    decided_constraint = require_contiguous  # type: ignore[assignment]
 
-                # for implicitly fallback ops, we conservatively requires
-                # contiguous input since some eager kernels does not
-                # support non-contiguous inputs. They may silently cause
-                # accuracy problems. Check https://github.com/pytorch/pytorch/issues/140452
                 make_fallback(target, layout_constraint=decided_constraint)
 
             elif get_decompositions([target]):
diff --git a/torch/_library/custom_ops.py b/torch/_library/custom_ops.py
@@ -605,7 +605,7 @@ def _register_to_dispatcher(self) -> None:
 
         lib.define(
             schema_str,
-            tags=[_C.Tag.pt2_compliant_tag, _C.Tag.needs_fixed_stride_order],
+            tags=[_C.Tag.pt2_compliant_tag, _C.Tag.needs_exact_strides],
         )
         self._opoverload = utils.lookup_op(self._qualname)
 
diff --git a/torch/_library/utils.py b/torch/_library/utils.py
@@ -267,6 +267,16 @@ def normalize_args_kwargs(
     return tuple(new_args), new_kwargs
 
 
+def needs_exact_strides(op: torch._ops.OpOverload):
+    if torch._C.Tag.needs_exact_strides in op.tags:
+        return True
+    if torch._C.Tag.flexible_layout in op.tags:
+        return False
+    if torch._C.Tag.needs_fixed_stride_order in op.tags:
+        return False
+    return not is_builtin(op)
+
+
 def hop_schema_from_fx_node(node):
     from torchgen.gen_schema_utils import FunctionSchemaGen
 
diff --git a/torch/fx/experimental/proxy_tensor.py b/torch/fx/experimental/proxy_tensor.py
@@ -1153,7 +1153,7 @@ def _should_save_arg_kwarg_vals(
                 f"propagate the FakeTensor vals. Please file an issue."
             )
     if isinstance(target, torch._ops.OpOverload):
-        return torch._C.Tag.needs_exact_strides in target.tags
+        return torch._library.utils.needs_exact_strides(target)
     return False
 
 
