[LT] Sync with the lazy_tensor_staging branch (#69527)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/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
diff --git a/tools/build_variables.bzl b/tools/build_variables.bzl
index 6239c6e..3e1f24b 100644
--- a/tools/build_variables.bzl
+++ b/tools/build_variables.bzl
@@ -374,13 +374,14 @@
"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 @@
"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",
]
diff --git a/tools/codegen/dest/lazy_ts_lowering.py b/tools/codegen/dest/lazy_ts_lowering.py
index 5e81746..32d505c 100644
--- a/tools/codegen/dest/lazy_ts_lowering.py
+++ b/tools/codegen/dest/lazy_ts_lowering.py
@@ -35,7 +35,7 @@
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
diff --git a/torch/csrc/lazy/core/internal_ops/ltc_ops.cpp b/torch/csrc/lazy/core/internal_ops/ltc_ops.cpp
deleted file mode 100644
index 0f9ccad..0000000
--- a/torch/csrc/lazy/core/internal_ops/ltc_ops.cpp
+++ /dev/null
@@ -1,30 +0,0 @@
-#include <torch/csrc/lazy/core/internal_ops/ltc_ops.h>
-
-namespace torch {
-namespace lazy {
-
-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
-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
diff --git a/torch/csrc/lazy/core/internal_ops/ltc_ops.h b/torch/csrc/lazy/core/internal_ops/ltc_ops.h
index f1f67ca..a6e1c5e 100644
--- a/torch/csrc/lazy/core/internal_ops/ltc_ops.h
+++ b/torch/csrc/lazy/core/internal_ops/ltc_ops.h
@@ -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 @@
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
diff --git a/torch/csrc/lazy/core/lazy_view.h b/torch/csrc/lazy/core/lazy_view.h
index aadd9f4..8116f6f 100644
--- a/torch/csrc/lazy/core/lazy_view.h
+++ b/torch/csrc/lazy/core/lazy_view.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 @@
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 @@
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 @@
int64_t dim2 = 1;
};
-struct ViewInfo {
+struct TORCH_API ViewInfo {
enum class Type {
kInvalid,
kNarrow,
@@ -96,7 +96,7 @@
// 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 @@
size_t generation_ = 0;
};
-class LazyView {
+class TORCH_API LazyView {
public:
LazyView(Shape shape, std::shared_ptr<Alias> alias, ViewInfo view_info);
LazyView(
diff --git a/torch/csrc/lazy/core/view_ops/as_strided.h b/torch/csrc/lazy/core/view_ops/as_strided.h
index 5dae891..e293f81 100644
--- a/torch/csrc/lazy/core/view_ops/as_strided.h
+++ b/torch/csrc/lazy/core/view_ops/as_strided.h
@@ -7,7 +7,7 @@
namespace torch {
namespace lazy {
-class AsStrided : public TsNode {
+class TORCH_API AsStrided : public TsNode {
public:
AsStrided(
const Value& input,
diff --git a/torch/csrc/lazy/core/view_ops/as_strided_view_update.h b/torch/csrc/lazy/core/view_ops/as_strided_view_update.h
index 27e4a93..85954ac 100644
--- a/torch/csrc/lazy/core/view_ops/as_strided_view_update.h
+++ b/torch/csrc/lazy/core/view_ops/as_strided_view_update.h
@@ -7,7 +7,7 @@
namespace torch {
namespace lazy {
-class AsStridedViewUpdate : public TsNode {
+class TORCH_API AsStridedViewUpdate : public TsNode {
public:
AsStridedViewUpdate(
const Value& target,
diff --git a/torch/csrc/lazy/core/view_ops/diagonal.h b/torch/csrc/lazy/core/view_ops/diagonal.h
index a55027b..2fdc6d3 100644
--- a/torch/csrc/lazy/core/view_ops/diagonal.h
+++ b/torch/csrc/lazy/core/view_ops/diagonal.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);
diff --git a/torch/csrc/lazy/core/view_ops/diagonal_view_update.h b/torch/csrc/lazy/core/view_ops/diagonal_view_update.h
index b4af441..813e2ae 100644
--- a/torch/csrc/lazy/core/view_ops/diagonal_view_update.h
+++ b/torch/csrc/lazy/core/view_ops/diagonal_view_update.h
@@ -5,7 +5,7 @@
namespace torch {
namespace lazy {
-class DiagonalViewUpdate : public TsNode {
+class TORCH_API DiagonalViewUpdate : public TsNode {
public:
DiagonalViewUpdate(
const Value& target,
diff --git a/torch/csrc/lazy/core/view_ops/narrow.h b/torch/csrc/lazy/core/view_ops/narrow.h
index d51a545..f017f34 100644
--- a/torch/csrc/lazy/core/view_ops/narrow.h
+++ b/torch/csrc/lazy/core/view_ops/narrow.h
@@ -5,7 +5,7 @@
namespace torch {
namespace lazy {
-class Narrow : public TsNode {
+class TORCH_API Narrow : public TsNode {
public:
Narrow(
const Value& input,
diff --git a/torch/csrc/lazy/core/view_ops/narrow_view_update.h b/torch/csrc/lazy/core/view_ops/narrow_view_update.h
index 44bb73f..f1d4cbc 100644
--- a/torch/csrc/lazy/core/view_ops/narrow_view_update.h
+++ b/torch/csrc/lazy/core/view_ops/narrow_view_update.h
@@ -5,7 +5,7 @@
namespace torch {
namespace lazy {
-class NarrowViewUpdate : public TsNode {
+class TORCH_API NarrowViewUpdate : public TsNode {
public:
NarrowViewUpdate(
const Value& input,
diff --git a/torch/csrc/lazy/core/view_ops/permute.h b/torch/csrc/lazy/core/view_ops/permute.h
index f034d92..9d84d0d 100644
--- a/torch/csrc/lazy/core/view_ops/permute.h
+++ b/torch/csrc/lazy/core/view_ops/permute.h
@@ -5,7 +5,7 @@
namespace torch {
namespace lazy {
-class Permute : public TsNode {
+class TORCH_API Permute : public TsNode {
public:
Permute(const Value& input, std::vector<int64_t> dims);
diff --git a/torch/csrc/lazy/core/view_ops/resize.h b/torch/csrc/lazy/core/view_ops/resize.h
index 0b59c4a..0da99d0 100644
--- a/torch/csrc/lazy/core/view_ops/resize.h
+++ b/torch/csrc/lazy/core/view_ops/resize.h
@@ -5,7 +5,7 @@
namespace torch {
namespace lazy {
-class Resize : public TsNode {
+class TORCH_API Resize : public TsNode {
public:
Resize(const Value& input, std::vector<int64_t> size);
diff --git a/torch/csrc/lazy/core/view_ops/select.h b/torch/csrc/lazy/core/view_ops/select.h
index dafce71..26cc7dd 100644
--- a/torch/csrc/lazy/core/view_ops/select.h
+++ b/torch/csrc/lazy/core/view_ops/select.h
@@ -5,7 +5,7 @@
namespace torch {
namespace lazy {
-class Select : public TsNode {
+class TORCH_API Select : public TsNode {
public:
Select(
const Value& input,
diff --git a/torch/csrc/lazy/core/view_ops/select_view_update.h b/torch/csrc/lazy/core/view_ops/select_view_update.h
index 907c7f0..6d33e1b 100644
--- a/torch/csrc/lazy/core/view_ops/select_view_update.h
+++ b/torch/csrc/lazy/core/view_ops/select_view_update.h
@@ -5,7 +5,7 @@
namespace torch {
namespace lazy {
-class SelectViewUpdate : public TsNode {
+class TORCH_API SelectViewUpdate : public TsNode {
public:
SelectViewUpdate(
const Value& target,
diff --git a/torch/csrc/lazy/core/view_ops/view.cpp b/torch/csrc/lazy/core/view_ops/view.cpp
index c76aca8..1a0e509 100644
--- a/torch/csrc/lazy/core/view_ops/view.cpp
+++ b/torch/csrc/lazy/core/view_ops/view.cpp
@@ -1,5 +1,7 @@
#include <torch/csrc/lazy/core/view_ops/view.h>
+#include <ATen/InferSize.h>
+
namespace torch {
namespace lazy {
diff --git a/torch/csrc/lazy/core/view_ops/view.h b/torch/csrc/lazy/core/view_ops/view.h
index a1bf458..8b69d5b 100644
--- a/torch/csrc/lazy/core/view_ops/view.h
+++ b/torch/csrc/lazy/core/view_ops/view.h
@@ -7,7 +7,7 @@
namespace torch {
namespace lazy {
-class View : public TsNode {
+class TORCH_API View : public TsNode {
public:
View(const Value& input, std::vector<int64_t> output_size);