Add C++ function for torch.distributed.tensor._op_schema.is_view_op (pytorch#161595)

This seems to have been an especially slow one because of the repeated pybind access (schema is a pybind, as is arguments, and then we hit each argument). It's still ~~1% of total benchmark runtime because of the repeated single pybind function call, but that's a lot better.

Differential Revision: [D81530095](https://our.internmc.facebook.com/intern/diff/D81530095)

Pull Request resolved: pytorch#161595
Approved by: https://github.com/ezyang, https://github.com/bdhirsh
ghstack dependencies: pytorch#161466, pytorch#161586, pytorch#161590, pytorch#161591

Author: swolchok

torch/_C/__init__.pyi.in

@@ -951,6 +951,7 @@ class FunctionSchema:
         is_vararg: _bool,
         is_varret: _bool,
     ) -> None: ...
+    def _is_view_op(self) -> _bool: ...
 
 class _UpgraderEntry:
     bumped_at_version: _int

torch/csrc/jit/python/init.cpp

@@ -1962,6 +1962,16 @@ void initJITBindings(PyObject* module) {
       .def_property_readonly("overload_name", &FunctionSchema::overload_name)
       .def_property_readonly("arguments", &FunctionSchema::arguments)
       .def_property_readonly("returns", &FunctionSchema::returns)
+      .def(
+          "_is_view_op",
+          [](const FunctionSchema& self) -> bool {
+            for (const auto& arg : self.arguments()) {
+              if (arg.alias_info() && !arg.alias_info()->isWrite()) {
+                return true;
+              }
+            }
+            return false;
+          })
       .def(
           "is_backward_compatible_with",
           // FunctionSchema::isBackwardCompatibleWith has an extra

torch/distributed/tensor/_op_schema.py

@@ -477,10 +477,7 @@ def is_out_variant_op(self) -> bool:
         return "out" in self.op._schema.overload_name
 
     def is_view_op(self) -> bool:
-        return any(
-            a.alias_info is not None and not a.alias_info.is_write
-            for a in self.op._schema.arguments
-        )
+        return self.op._schema._is_view_op()
 
     def _recompute_comparison_key(self):
         if not self.schema_info: