Inductor respects exact strides on custom ops by default (pytorch#150511)

If a tag is not specified on a custom operator, then inductor will
assume that it needs exact strides.

Test Plan:
- tests + CI

Pull Request resolved: pytorch#150511
Approved by: https://github.com/eellison, https://github.com/shunting314
ghstack dependencies: pytorch#148104

Author: zou3519

test/inductor/test_triton_kernels.py

@@ -3451,7 +3451,6 @@ def impl2(x):
 
             lib.define(
                 "add_op(Tensor x, Tensor y) -> Tensor",
-                tags=[torch._C.Tag.needs_exact_strides],
             )
 
             def impl(x, y):
@@ -3465,7 +3464,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):

test/test_custom_ops.py

@@ -3838,14 +3838,15 @@ def vmap(info, in_dims, w, x=2, *, y=3, z):
             self.assertEqual(result, w * 2 * 3 * 42)
 
     def test_layout_constraint_tags(self):
+        needs_exact_strides = torch._C.Tag.needs_exact_strides
         needs_fixed_stride_order = torch._C.Tag.needs_fixed_stride_order
         flexible_layout = torch._C.Tag.flexible_layout
         # (tags, the result of the tag inference)
         tests = [
             ({needs_fixed_stride_order}, needs_fixed_stride_order),
             ({flexible_layout}, flexible_layout),
             # If no tags are provided, then the following is the default
-            (set(), needs_fixed_stride_order),
+            (set(), needs_exact_strides),
             # If multiple tags are provided, then we use the most constrained tag.
             ({flexible_layout, needs_fixed_stride_order}, needs_fixed_stride_order),
         ]

torch/_functorch/config.py

@@ -186,7 +186,7 @@ def remote_autograd_cache_default() -> Optional[bool]:
 # ProxyTensor tracing.
 custom_op_default_layout_constraint: Literal[
     "needs_exact_strides", "needs_fixed_stride_order", "flexible_layout"
-] = "needs_fixed_stride_order"
+] = "needs_exact_strides"
 
 
 # Run aot eager decomp partition with CrossRefFakeMode

torch/_library/custom_ops.py

@@ -615,7 +615,7 @@ def _register_to_dispatcher(self, tags: Sequence[_C.Tag]) -> None:
 
         lib.define(
             schema_str,
-            tags=[_C.Tag.pt2_compliant_tag, _C.Tag.needs_fixed_stride_order, *tags],
+            tags=[_C.Tag.pt2_compliant_tag, *tags],
         )
         self._opoverload = utils.lookup_op(self._qualname)