Codegen Non-Native IR Nodes (pytorch#76535)

Add codegen infrastructure to generate IR nodes for non-native ops.

The proposed change is to add a `non_native` key to the `{backend}_native_functions.yaml` file that contains schema definitions similar to what is found in `native_functions.yaml`. e.g.
```
non_native:
    ...
    - func: expand(Tensor input, int[] size, bool is_scalar_expand) -> Tensor
    ...
```
these definitions are parsed into a `LazyIrSchema` that can be used for generating IR nodes using `GenLazyIR`.

Fixes pytorch#74628

CC: @wconstab @desertfire @henrytwo

Pull Request resolved: pytorch#76535
Approved by: https://github.com/wconstab

Author: antoniojkim

BUILD.bazel

@@ -1865,6 +1865,7 @@ test_suite(
         "aten/src/ATen/templates/DispatchKeyNativeFunctions.cpp",
         "aten/src/ATen/templates/DispatchKeyNativeFunctions.h",
         "aten/src/ATen/templates/LazyIr.h",
+        "aten/src/ATen/templates/LazyNonNativeIr.h",
         "aten/src/ATen/templates/RegisterDispatchKey.cpp",
         "aten/src/ATen/native/native_functions.yaml",
         "aten/src/ATen/native/tags.yaml",

aten/src/ATen/native/ts_native_functions.yaml

@@ -178,3 +178,41 @@ supported:
   - _unsafe_view
 autograd:
   - max_pool3d
+
+# Ops that don't have a native schema definitions and are dispatched within Lazy Tensor Core
+non_native:
+  - func: scalar(Scalar value, ScalarType type) -> Tensor
+    opkind: at::prim::Constant
+    properties:
+      - ShapeCompute
+      - TreatScalarsAsConstants
+  - func: expand(Tensor input, int[] size, bool is_scalar_expand) -> Tensor
+  - func: view(Tensor input, int[] output_size) -> Tensor
+    properties:
+      - ShapeCompute
+  - func: cast(Tensor input, ScalarType dtype, ScalarType? stype) -> Tensor
+    opkind: ltc_cast
+    properties:
+      - ShapeCompute
+
+  # View ops only required until proper functionalization pass is introduced into LTC
+  - func: as_strided_view_update(Tensor target, Tensor input, int[] size, int[] stride, int storage_offset) -> Tensor
+    opkind: ltc_as_strided_view_update
+  - func: as_strided(Tensor input, int[] size, int[] stride, int storage_offset) -> Tensor
+  - func: diagonal_view_update(Tensor target, Tensor input, int offset, int dim1, int dim2) -> Tensor
+    opkind: ltc_diagonal_view_update
+    properties:
+      - ShapeCompute
+  - func: diagonal(Tensor input, int offset, int dim1, int dim2) -> Tensor
+  - func: narrow_view_update(Tensor input, Tensor source, int[] base_indices) -> Tensor
+    opkind: ltc_narrow_view_update
+  - func: narrow(Tensor input, int[] base_indices, int[] sizes) -> Tensor
+  - func: permute(Tensor input, int[] dims) -> Tensor
+  - func: resize(Tensor input, int[] size) -> Tensor
+  - func: select_view_update(Tensor target, Tensor source, int dim, int start, int end, int stride) -> Tensor
+    opkind: ltc_select_view_update
+    properties:
+      - ShapeCompute
+  - func: select(Tensor input, int dim, int start, int end, int stride) -> Tensor
+  - func: squeeze(Tensor input, int dim) -> Tensor
+  - func: unsqueeze(Tensor input, int dim) -> Tensor

aten/src/ATen/templates/LazyNonNativeIr.h

@@ -0,0 +1,11 @@
+#pragma once
+
+${lazy_non_native_ir_inc}
+
+// This file contains autogenerated LazyTensor Non Native IR nodes
+
+${namespace_prologue}
+
+${non_native_ir_nodes}
+
+${namespace_epilogue}

build.bzl

@@ -28,6 +28,7 @@ def define_targets(rules):
             ":DispatchKeyNativeFunctions.cpp",
             ":DispatchKeyNativeFunctions.h",
             ":LazyIr.h",
+            ":LazyNonNativeIr.h",
             ":RegisterDispatchKey.cpp",
             ":native_functions.yaml",
             ":shape_inference.h",
@@ -88,6 +89,7 @@ GENERATED_TESTING_PY = [
 
 GENERATED_LAZY_H = [
     "torch/csrc/lazy/generated/LazyIr.h",
+    "torch/csrc/lazy/generated/LazyNonNativeIr.h",
     "torch/csrc/lazy/generated/LazyNativeFunctions.h",
 ]
 

caffe2/CMakeLists.txt

@@ -380,6 +380,7 @@ if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE)
     list(APPEND GENERATED_H_TORCH
       "${TORCH_SRC_DIR}/csrc/autograd/generated/VariableType.h"
       "${TORCH_SRC_DIR}/csrc/lazy/generated/LazyIr.h"
+      "${TORCH_SRC_DIR}/csrc/lazy/generated/LazyNonNativeIr.h"
       "${TORCH_SRC_DIR}/csrc/lazy/generated/LazyNativeFunctions.h"
     )
   endif()
@@ -444,6 +445,7 @@ if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE)
     "${TORCH_ROOT}/aten/src/ATen/templates/DispatchKeyNativeFunctions.h"
     "${TORCH_ROOT}/aten/src/ATen/templates/DispatchKeyNativeFunctions.cpp"
     "${TORCH_ROOT}/aten/src/ATen/templates/LazyIr.h"
+    "${TORCH_ROOT}/aten/src/ATen/templates/LazyNonNativeIr.h"
     "${TORCH_ROOT}/aten/src/ATen/templates/RegisterDispatchKey.cpp"
     "${TOOLS_PATH}/autograd/templates/VariableType.h"
     "${TOOLS_PATH}/autograd/templates/VariableType.cpp"

tools/build_variables.bzl

@@ -417,35 +417,19 @@ lazy_tensor_core_sources = [
 # We can't build all of the ts backend under certain build configurations, e.g. mobile,
 # since it depends on things like autograd, meta functions, which may be disabled
 lazy_tensor_ts_sources = [
-    "torch/csrc/lazy/ts_backend/config.cpp",
     "torch/csrc/lazy/ts_backend/dynamic_ir.cpp",
+    "torch/csrc/lazy/ts_backend/config.cpp",
     "torch/csrc/lazy/ts_backend/ops/batch_norm_ops.cpp",
-    "torch/csrc/lazy/ts_backend/ops/random_ops.cpp",
-    "torch/csrc/lazy/ts_backend/ops/cast.cpp",
     "torch/csrc/lazy/ts_backend/ops/device_data.cpp",
-    "torch/csrc/lazy/ts_backend/ops/expand.cpp",
+    "torch/csrc/lazy/ts_backend/ops/random_ops.cpp",
     "torch/csrc/lazy/ts_backend/ops/generic.cpp",
-    "torch/csrc/lazy/ts_backend/ops/scalar.cpp",
-    "torch/csrc/lazy/ts_backend/view_ops/as_strided.cpp",
-    "torch/csrc/lazy/ts_backend/view_ops/as_strided_view_update.cpp",
-    "torch/csrc/lazy/ts_backend/view_ops/diagonal.cpp",
-    "torch/csrc/lazy/ts_backend/view_ops/diagonal_view_update.cpp",
-    "torch/csrc/lazy/ts_backend/view_ops/narrow.cpp",
-    "torch/csrc/lazy/ts_backend/view_ops/narrow_view_update.cpp",
-    "torch/csrc/lazy/ts_backend/view_ops/permute.cpp",
-    "torch/csrc/lazy/ts_backend/view_ops/resize.cpp",
-    "torch/csrc/lazy/ts_backend/view_ops/select.cpp",
-    "torch/csrc/lazy/ts_backend/view_ops/squeeze.cpp",
-    "torch/csrc/lazy/ts_backend/view_ops/unsqueeze.cpp",
-    "torch/csrc/lazy/ts_backend/view_ops/select_view_update.cpp",
-    "torch/csrc/lazy/ts_backend/view_ops/view.cpp",
-    "torch/csrc/lazy/ts_backend/ts_node.cpp",
     "torch/csrc/lazy/ts_backend/tensor_aten_ops.cpp",
     "torch/csrc/lazy/ts_backend/ts_autograd_functions.cpp",
     "torch/csrc/lazy/ts_backend/ts_backend_impl.cpp",
     "torch/csrc/lazy/ts_backend/ts_eager_fallback.cpp",
     "torch/csrc/lazy/ts_backend/ts_lowering_context.cpp",
     "torch/csrc/lazy/ts_backend/ts_native_functions.cpp",
+    "torch/csrc/lazy/ts_backend/ts_node.cpp",
     "torch/csrc/lazy/ts_backend/ts_node_lowering.cpp",
 ]
 

tools/test/test_gen_backend_stubs.py

@@ -237,7 +237,7 @@ def test_unrecognized_key(self) -> None:
         output_error = self.get_errors_from_gen_backend_stubs(yaml_str)
         self.assertExpectedInline(
             output_error,
-            """ contains unexpected keys: invalid_key. Only the following keys are supported: backend, class_name, cpp_namespace, extra_headers, supported, autograd, full_codegen""",  # noqa: B950
+            """ contains unexpected keys: invalid_key. Only the following keys are supported: backend, class_name, cpp_namespace, extra_headers, supported, autograd, full_codegen, non_native""",  # noqa: B950
         )
 
     # if use_out_as_primary is provided, it must be a bool

torch/csrc/lazy/core/ir.cpp

@@ -19,6 +19,10 @@ hash_t Output::hash() const {
   return HashCombine(node->hash(), Hash(index));
 }
 
+hash_t Output::shapeHash() const {
+  return HashCombine(node->shapeHash(), Hash(index));
+}
+
 std::string Output::ToString() const {
   std::stringstream ss;
   ss << node->ToString() << ", index=" << index;
@@ -144,7 +148,7 @@ std::string Node::ToString() const {
 
 void Node::AddOperand(NodePtr node, size_t index) {
   CHECK_LT(index, node->num_outputs());
-  operands_.push_back(std::move(node));
+  operands_.push_back(node);
   operands_as_outputs_.emplace_back(operands_.back().get(), index);
 }
 

torch/csrc/lazy/core/ir.h

@@ -214,6 +214,7 @@ struct TORCH_API Output {
       : node(node), index(index) {}
 
   hash_t hash() const;
+  hash_t shapeHash() const;
 
   bool operator==(const Output& rhs) const {
     return node == rhs.node && index == rhs.index;

torch/csrc/lazy/core/ops/utils.h

@@ -6,33 +6,33 @@
 namespace torch {
 namespace lazy {
 
-bool StrideIsSupported(c10::ArrayRef<int64_t> stride);
+TORCH_API bool StrideIsSupported(c10::ArrayRef<int64_t> stride);
 
-std::vector<int64_t> GetArrayStridePermutation(c10::ArrayRef<int64_t> stride);
+TORCH_API std::vector<int64_t> GetArrayStridePermutation(c10::ArrayRef<int64_t> stride);
 
-Shape MakeDiagonalShape(
+TORCH_API Shape MakeDiagonalShape(
     const Shape& shape,
     int64_t offset,
     int64_t dim1,
     int64_t dim2);
 
-Shape MakePermuteShape(
+TORCH_API Shape MakePermuteShape(
     const Shape& source_shape,
     c10::ArrayRef<int64_t> permutation);
 
-Shape MakeSelectShape(
+TORCH_API Shape MakeSelectShape(
     const Shape& shape,
     int64_t dim,
     int64_t start,
     int64_t end,
     int64_t stride);
 
-int64_t GetStride(int64_t start, int64_t end, int64_t stride);
+TORCH_API int64_t GetStride(int64_t start, int64_t end, int64_t stride);
 
-std::vector<int64_t> BuildSqueezedDimensions(c10::ArrayRef<int64_t> dimensions,
+TORCH_API std::vector<int64_t> BuildSqueezedDimensions(c10::ArrayRef<int64_t> dimensions,
                                              int64_t squeeze_dim);
 
-std::vector<int64_t> BuildUnsqueezedDimensions(
+TORCH_API std::vector<int64_t> BuildUnsqueezedDimensions(
     c10::ArrayRef<int64_t> dimensions,
     int64_t squeeze_dim);