[LT] Sync with the lazy_tensor_staging branch (pytorch#69527)

Summary:
Pull Request resolved: pytorch#69527

- Add missing TORCH_API in class/struct declarations;
- Fix internal op declarations in ltc_ops;
- Update lazy_ts_lowering.py

Test Plan: Imported from OSS

Reviewed By: alanwaketan

Differential Revision: D32918929

Pulled By: desertfire

fbshipit-source-id: e956d51aff5ef593fdf4cd5ad2a38e38788913d8

Author: desertfire

tools/build_variables.bzl

@@ -374,13 +374,14 @@ lazy_tensor_core_sources = [
     "torch/csrc/lazy/core/config.cpp",
     "torch/csrc/lazy/core/hash.cpp",
     "torch/csrc/lazy/core/helpers.cpp",
-    "torch/csrc/lazy/core/internal_ops/ltc_ops.cpp",
     "torch/csrc/lazy/core/ir.cpp",
     "torch/csrc/lazy/core/ir_dump_util.cpp",
     "torch/csrc/lazy/core/ir_metadata.cpp",
     "torch/csrc/lazy/core/ir_util.cpp",
+    "torch/csrc/lazy/core/lazy_view.cpp",
     "torch/csrc/lazy/core/permutation_util.cpp",
     "torch/csrc/lazy/core/shape.cpp",
+    "torch/csrc/lazy/core/tensor_util.cpp",
     "torch/csrc/lazy/core/view_ops/as_strided.cpp",
     "torch/csrc/lazy/core/view_ops/as_strided_view_update.cpp",
     "torch/csrc/lazy/core/view_ops/diagonal.cpp",
@@ -392,8 +393,6 @@ lazy_tensor_core_sources = [
     "torch/csrc/lazy/core/view_ops/select.cpp",
     "torch/csrc/lazy/core/view_ops/select_view_update.cpp",
     "torch/csrc/lazy/core/view_ops/view.cpp",
-    "torch/csrc/lazy/core/tensor_util.cpp",
-    "torch/csrc/lazy/core/lazy_view.cpp",
     "torch/csrc/lazy/ts_backend/config.cpp",
     "torch/csrc/lazy/ts_backend/ts_node.cpp",
 ]

tools/codegen/dest/lazy_ts_lowering.py

@@ -35,7 +35,7 @@ def ts_lowering_body(f: Union[NativeFunctionsGroup, NativeFunction]) -> str:
     size_t i = 0;
     {emplace_arguments_str}
     {emplace_kwarguments}
-    TSOpVector {schema.aten_name}_out = compiler::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    torch::lazy::TSOpVector {schema.aten_name}_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
     CHECK_EQ({schema.aten_name}_out.size(), {len(func.returns)});
 
     // TODO: need to call GenerateClone sometimes? Or else return LowerBuiltIn() directly

torch/csrc/lazy/core/internal_ops/ltc_ops.cpp

@@ -8,7 +8,7 @@
 namespace torch {
 namespace lazy {
 
-class OpKindWrapper {
+class TORCH_API OpKindWrapper {
  public:
   explicit OpKindWrapper(const char* name) : name_(name) {}
 
@@ -31,25 +31,28 @@ class OpKindWrapper {
   mutable std::once_flag once_;
 };
 
-// TODO: some ops may be obsolete
-extern const OpKindWrapper ltc_all_to_all;
-extern const OpKindWrapper ltc_cast;
-extern const OpKindWrapper ltc_collective_permute;
-extern const OpKindWrapper ltc_cross_replica_sum;
-extern const OpKindWrapper ltc_device_data;
-extern const OpKindWrapper ltc_get_dimensions_size;
-extern const OpKindWrapper ltc_moving_average;
-extern const OpKindWrapper ltc_nms;
-extern const OpKindWrapper ltc_not_supported;
-extern const OpKindWrapper ltc_replication_pad;
-extern const OpKindWrapper ltc_replication_pad_backward;
-extern const OpKindWrapper ltc_tensor_data;
+const OpKindWrapper ltc_all_to_all("lazy_tensors::all_to_all");
+const OpKindWrapper ltc_cast("lazy_tensors::cast");
+const OpKindWrapper ltc_collective_permute("lazy_tensors::collective_permute");
+const OpKindWrapper ltc_cross_replica_sum("lazy_tensors::cross_replica_sum");
+const OpKindWrapper ltc_device_data("lazy_tensors::device_data");
+const OpKindWrapper ltc_get_dimensions_size(
+    "lazy_tensors::ltc_get_dimensions_size");
+const OpKindWrapper ltc_moving_average("lazy_tensors::moving_average");
+const OpKindWrapper ltc_nms("lazy_tensors::nms");
+const OpKindWrapper ltc_not_supported("lazy_tensors::not_supported");
+const OpKindWrapper ltc_replication_pad("lazy_tensors::replication_pad");
+const OpKindWrapper ltc_replication_pad_backward(
+    "lazy_tensors::replication_pad_backward");
+const OpKindWrapper ltc_tensor_data("lazy_tensors::tensor_data");
 
 // For view ops
-extern const OpKindWrapper ltc_as_strided_view_update;
-extern const OpKindWrapper ltc_diagonal_view_update;
-extern const OpKindWrapper ltc_narrow_view_update;
-extern const OpKindWrapper ltc_select_view_update;
+const OpKindWrapper ltc_as_strided_view_update(
+    "lazy_tensors::as_strided_view_update");
+const OpKindWrapper ltc_diagonal_view_update(
+    "lazy_tensors::diagonal_view_update");
+const OpKindWrapper ltc_narrow_view_update("lazy_tensors::narrow_view_update");
+const OpKindWrapper ltc_select_view_update("lazy_tensors::select_view_update");
 
 } // namespace lazy
 } // namespace torch

torch/csrc/lazy/core/internal_ops/ltc_ops.h

@@ -10,7 +10,7 @@
 namespace torch {
 namespace lazy {
 
-struct SelectInfo {
+struct TORCH_API SelectInfo {
   bool operator==(const SelectInfo& ref) const {
     return dim == ref.dim && start == ref.start && end == ref.end &&
         stride == ref.stride;
@@ -22,7 +22,7 @@ struct SelectInfo {
   int64_t stride = 0;
 };
 
-struct AsStridedInfo {
+struct TORCH_API AsStridedInfo {
   bool operator==(const AsStridedInfo& ref) const {
     return offset == ref.offset && stride == ref.stride;
   }
@@ -31,7 +31,7 @@ struct AsStridedInfo {
   int64_t offset = 0;
 };
 
-struct DiagonalInfo {
+struct TORCH_API DiagonalInfo {
   bool operator==(const DiagonalInfo& ref) const {
     return offset == ref.offset && dim1 == ref.dim1 && dim2 == ref.dim2;
   }
@@ -41,7 +41,7 @@ struct DiagonalInfo {
   int64_t dim2 = 1;
 };
 
-struct ViewInfo {
+struct TORCH_API ViewInfo {
   enum class Type {
     kInvalid,
     kNarrow,
@@ -96,7 +96,7 @@ struct ViewInfo {
 
 // When a "view" (capture by reference) is taken on a node, an Alias object is
 // created on the captured node itself, with its current IR Node value.
-class Alias {
+class TORCH_API Alias {
  public:
   struct UpdateData {
     Value ir_value;
@@ -127,7 +127,7 @@ class Alias {
   size_t generation_ = 0;
 };
 
-class LazyView {
+class TORCH_API LazyView {
  public:
   LazyView(Shape shape, std::shared_ptr<Alias> alias, ViewInfo view_info);
   LazyView(

torch/csrc/lazy/core/lazy_view.h

@@ -7,7 +7,7 @@
 namespace torch {
 namespace lazy {
 
-class AsStrided : public TsNode {
+class TORCH_API AsStrided : public TsNode {
  public:
   AsStrided(
       const Value& input,

torch/csrc/lazy/core/view_ops/as_strided.h

@@ -7,7 +7,7 @@
 namespace torch {
 namespace lazy {
 
-class AsStridedViewUpdate : public TsNode {
+class TORCH_API AsStridedViewUpdate : public TsNode {
  public:
   AsStridedViewUpdate(
       const Value& target,

torch/csrc/lazy/core/view_ops/as_strided_view_update.h

@@ -5,7 +5,7 @@
 namespace torch {
 namespace lazy {
 
-class Diagonal : public TsNode {
+class TORCH_API Diagonal : public TsNode {
  public:
   Diagonal(const Value& input, int64_t offset, int64_t dim1, int64_t dim2);
 

torch/csrc/lazy/core/view_ops/diagonal.h

@@ -5,7 +5,7 @@
 namespace torch {
 namespace lazy {
 
-class DiagonalViewUpdate : public TsNode {
+class TORCH_API DiagonalViewUpdate : public TsNode {
  public:
   DiagonalViewUpdate(
       const Value& target,

torch/csrc/lazy/core/view_ops/diagonal_view_update.h

@@ -5,7 +5,7 @@
 namespace torch {
 namespace lazy {
 
-class Narrow : public TsNode {
+class TORCH_API Narrow : public TsNode {
  public:
   Narrow(
       const Value& input,