[POC] Dynamic `expand` with `SymInt` implementation #3558

miladm · 2022-05-10T04:03:04Z

This is the POC implementation of torch.Tensor.expand op based on the PyTorch SymInt POC implementation PR.

Action items to unblock:

torch_xla/csrc/aten_xla_type.cpp

JackCaoG · 2022-05-10T17:42:37Z

Does this pr build locally on your end? Build on CI failed with conflicts.

torch_xla/csrc/aten_xla_type.cpp

JackCaoG · 2022-05-10T17:56:30Z

torch_xla/csrc/aten_xla_type.cpp

+  for (int index = 0; i < sizes.size(); i++) {
+    auto _symbolicIntNode = sizes[i].toSymbolicIntNode();
+    auto _sizenode = MakeNode<SizeNode>(_symbolicIntNode);
+    upper_bound.push_back(_sizenode.getStaticValue());
+    dynamic_dims.push_back(_sizenode.isDynamic());
+  }


we can probably put these in a helper function for now. I felt like we should be able to codegen these in the future too.

JackCaoG · 2022-05-10T18:00:23Z

torch_xla/csrc/ops/ops.cpp

+
+    /* Construct Upper Bound Tensor Shape */
+    xla::XlaOp upper_bound_size_input =
+        xla::Parameter(loctx->builder(), 0, target_shape, "upper_bound_size");


hmm, it is rare to use xla::Parameter in the lowering, you can use xla::Zero and xla::Broadcast to achieve the same.

Why do I use to two API calls to reach the same goal instead of one?

torch_xla/csrc/data_ops.cpp

torch_xla/csrc/aten_xla_type.cpp

miladm · 2022-05-10T18:39:48Z

This PR doesn't build at the moment because the upstream layer LTC doesn't yet have API support for expand with SymInt. @JackCaoG @Krovatkin

torch_xla/csrc/ops/ops.cpp

miladm · 2022-05-16T05:40:24Z

Update: The current unit test checks the expand.SymInt code path. It does not check the dynamic dimension propagation across a SymInt op since DimensionNode::isDynamic implementation is currently WIP.

CC @JackCaoG @Krovatkin

miladm · 2022-05-16T05:45:16Z

torch_xla/csrc/aten_xla_type.cpp

+  std::vector<torch::lazy::NodePtr> size_nodes;
+  std::vector<int64_t> upper_bounds;
+  std::vector<bool> dynamic_dims;
+  /* TODO: move this code to a helper function */


Considering to move this code to a helper function. My candidates are creating a new helper file, helpers.h, or tensor_util.h? @JackCaoG wdyt?

you could also consider putting a helper in dynamic_ir.h. And I can make a copy for the TS backend. That's a disadvantage of code duplication. Hopefully, we won't need to do it often.

Krovatkin · 2022-05-16T16:12:22Z

torch_xla/csrc/aten_xla_type.cpp

+  std::vector<torch::lazy::NodePtr> size_nodes;
+  std::vector<int64_t> upper_bounds;
+  std::vector<bool> dynamic_dims;
+  /* TODO: move this code to a helper function */


you could also consider putting a helper in dynamic_ir.h. And I can make a copy for the TS backend. That's a disadvantage of code duplication. Hopefully, we won't need to do it often.

torch_xla/csrc/aten_xla_type.cpp

Krovatkin · 2022-05-16T16:38:27Z

torch_xla/csrc/data_ops.cpp

+                        absl::Span<const xla::XlaOp> output_sizes) {
+  for (int i = 0; i < output_sizes.size(); i++) {
+    xla::Shape dim_shape = XlaHelpers::ShapeOfXlaOp(output_sizes[i]);
+    if (dim_shape.is_dynamic()) {


and if the shape isn't dynamic we will be using static dimensions from InferShape? I'm asking because we would like shapes to appear symbolically in a graph and I'm not sure this would happen if dimensions are static? @miladm @JackCaoG

Yes, dimensions from InferShape set shape. This means when a dimension is static, the upper_bound values are expected to be equal to the "true" dimensions. Does this align with your side @Krovatkin?

Can you elaborate on: we would like shapes to appear symbolically in a graph and I'm not sure this would happen if dimensions are static? @Krovatkin

Krovatkin · 2022-05-16T17:02:54Z

xla_native_functions.yaml

@@ -116,6 +116,7 @@ supported:
  - erfinv
  - exp
  - expand
+  - expand.SymInt


I'm actually not seeing the exact place where your hooking in XLATensor::expand into a dispatcher? is the plan to use codegen for that?

Our current codegen creates XLANativeFunctions.h for aten_xla_type.cpp. In it, you find the following definitions. Does it answer your question @Krovatkin?

static at::Tensor expand(const at::Tensor & self, at::IntArrayRef size, bool implicit); static at::Tensor expand(const at::Tensor & self, c10::SymIntArrayRef size, bool implicit);

JackCaoG

still WIP for the review..

JackCaoG · 2022-05-16T22:35:45Z

test/cpp/run_tests.sh

-  -DPYTHON_INCLUDE_DIR=$(python -c "from distutils.sysconfig import get_python_inc; print(get_python_inc())") \
-  -DPYTHON_LIBRARY=$(python -c "import distutils.sysconfig as sysconfig; print(sysconfig.get_config_var('LIBDIR') + '/' + sysconfig.get_config_var('LDLIBRARY'))")
+  -DPYTHON_INCLUDE_DIR=$(python3 -c "from distutils.sysconfig import get_python_inc; print(get_python_inc())") \
+  -DPYTHON_LIBRARY=$(python3 -c "import distutils.sysconfig as sysconfig; print(sysconfig.get_config_var('LIBDIR') + '/' + sysconfig.get_config_var('LDLIBRARY'))")


There was an issue by forcing python to python3, @yeounoh might have more context. If this is for the test purpose you can use

sudo update-alternatives --install /usr/bin/python python /usr/bin/python3.8 100

on your tpuvm

test/cpp/test_aten_xla_tensor.cpp

torch_xla/csrc/data_ops.cpp

torch_xla/csrc/ops/dynamic_ir.cpp

JackCaoG · 2022-05-16T23:19:19Z

torch_xla/csrc/ops/dynamic_ir.cpp

+
+std::string SizeNode::ToString() const { return "SizeNode"; }
+
+SizeAdd::SizeAdd(XlaValue a, XlaValue b)


should we check that XlaValue actually contains DimensionNode and store them as DimensionNode diretly? This can save us from doing multiple dynamic_cast on getStaticValue.

I am flexible. @Krovatkin wdyt?

JackCaoG · 2022-05-16T23:20:15Z

torch_xla/csrc/ops/dynamic_ir.h

+  SizeAdd(XlaValue a, XlaValue b);
+  int64_t getStaticValue() const override;
+  std::string ToString() const override;
+};


why these operator does not have Lower?

Not sure if we need to lower these ops as they return upper bound information. Unless we directly pass SizeAdd to downstream ops like expand. @Krovatkin wdyt?

To help the discussion I will push a lowering implementation shortly. Let's continue discussing.

JackCaoG · 2022-05-16T23:21:03Z

torch_xla/csrc/ops/dynamic_ir.cpp

+SizeDiv::SizeDiv(XlaValue a, XlaValue b)
+    : DimensionNode(
+          torch::lazy::OpKind{c10::Symbol::fromQualString("aten::div")}, {a, b},
+          torch::lazy::MHash(1)){};


I take this hash is a hack for now..

torch_xla/csrc/ops/ops.cpp

torch_xla/csrc/tensor_methods.cpp

miladm · 2022-05-19T01:18:50Z

test/cpp/test_aten_xla_tensor.cpp

+    c10::SymInt xla_y0_size = xla_y.sym_sizes()[0];
+    torch::Tensor xla_a = CopyToDevice(a, device);
+    torch::Tensor xla_b = xla_a.expand(
+        c10::SymIntArrayRef({xla_y0_size, c10::SymInt(3), c10::SymInt(4)}),


@Krovatkin I hope constructing c10::SymIntArrayRef makes sense (ref). Let me know if it should be done differently.

it looks good. I think just plain xla_a.expand_symint({xla_y0_size, c10::SymInt(3), c10::SymInt(4)}, .... should've worked. If not we should look into that.

@Krovatkin as discussed this morning, looks like this line causes the failure raised here. What's your guidance?

…temp measure for POC

…77917/files

miladm added this to the Dynamic Shape milestone May 10, 2022

miladm added the DO_NOT_MERGE_YET For PRs which cannot be merged, despite tests passing label May 10, 2022

miladm self-assigned this May 10, 2022

miladm requested a review from JackCaoG May 10, 2022 04:06

miladm changed the title ~~dynamic expand with symint implementation POC~~ [POC] Dynamic expand with SymInt implementation May 10, 2022

miladm requested review from Krovatkin and ezyang May 10, 2022 07:01

miladm commented May 10, 2022

View reviewed changes

torch_xla/csrc/aten_xla_type.cpp Outdated Show resolved Hide resolved

JackCaoG reviewed May 10, 2022

View reviewed changes

miladm requested a review from wconstab May 10, 2022 18:53

miladm marked this pull request as draft May 10, 2022 18:58

JackCaoG reviewed May 10, 2022

View reviewed changes

torch_xla/csrc/ops/ops.cpp Outdated Show resolved Hide resolved

JackCaoG reviewed May 10, 2022

View reviewed changes

torch_xla/csrc/ops/ops.cpp Outdated Show resolved Hide resolved

miladm force-pushed the expand_symint_poc branch from d07f724 to 2d2025f Compare May 15, 2022 22:31

miladm commented May 16, 2022

View reviewed changes

Krovatkin reviewed May 16, 2022

View reviewed changes

JackCaoG reviewed May 16, 2022

View reviewed changes

JackCaoG reviewed May 17, 2022

View reviewed changes

torch_xla/csrc/tensor_methods.cpp Outdated Show resolved Hide resolved

miladm changed the title ~~[POC] Dynamic expand with SymInt implementation~~ [POC][WIP] Dynamic expand with SymInt implementation May 17, 2022

miladm commented May 19, 2022

View reviewed changes

miladm added BLOCKED dynamism Dynamic Shape Features labels May 21, 2022

miladm mentioned this pull request May 21, 2022

[POC] Support of dynamic shapes for the fill_ op #3100

Closed

miladm added 22 commits July 4, 2022 06:32

test fix

4496a2f

lint

7fa7354

transitioned to ExpandDynamic class

04367ee

addressed some of the review feedbacks

7cedbaa

addressed some of the review feedbacks

0a62719

made a few improvements

7389491

instruced SymIntElements struct

b649fcf

linter

9f16cce

linger

22326d3

cleanup

a05164b

improved test code - test is yet to be run

f63a5ba

adding lowering for arithmetic size ops

f34b4b5

linter

c849434

enabling dynamic shape vs. nonzero by commenting the fallback code - …

b2c4366

…temp measure for POC

added support for DimensionNode to inherit from LTC DimensionNode

a8a2e84

linter

4f32b14

added isDynamic logic - ref: https://github.com/pytorch/pytorch/pull/…

e76134f

…77917/files

linter

d532b84

upgrade to latest LTC APIs

021046c

linter

051e877

linter

d9e19d0

updated expand.SymInt API in PT/XLA

5681853

miladm force-pushed the expand_symint_poc branch from 1961ce6 to 5681853 Compare July 4, 2022 06:32

miladm added 5 commits July 4, 2022 06:33

linter

ba1496b

handle non-dynamic symint calls, fix xlaop for shapes logic

b754eac

addressed conflicts

e7fcf74

linter

cdb7b74

added more test flags

cc59c69

miladm changed the title ~~[POC][WIP] Dynamic expand with SymInt implementation~~ [POC] Dynamic expand with SymInt implementation Jul 12, 2022

miladm mentioned this pull request Jul 12, 2022

Dynamic expand with SymInt implementation #3702

Draft

1 task

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[POC] Dynamic `expand` with `SymInt` implementation #3558

[POC] Dynamic `expand` with `SymInt` implementation #3558

miladm commented May 10, 2022 •

edited

JackCaoG commented May 10, 2022

JackCaoG May 10, 2022

JackCaoG May 10, 2022

miladm May 10, 2022 •

edited

miladm commented May 10, 2022

miladm commented May 16, 2022

miladm May 16, 2022

Krovatkin May 16, 2022

Krovatkin May 16, 2022

Krovatkin May 16, 2022

miladm May 16, 2022 •

edited

Krovatkin May 16, 2022

miladm May 16, 2022 •

edited

JackCaoG left a comment •

edited

JackCaoG May 16, 2022

JackCaoG May 16, 2022

miladm May 19, 2022

JackCaoG May 16, 2022

miladm May 19, 2022 •

edited

JackCaoG May 16, 2022

miladm May 19, 2022 •

edited

Krovatkin Jul 7, 2022

miladm Jul 7, 2022 •

edited


		std::string SizeNode::ToString() const { return "SizeNode"; }

		SizeAdd::SizeAdd(XlaValue a, XlaValue b)

[POC] Dynamic expand with SymInt implementation #3558

Are you sure you want to change the base?

[POC] Dynamic expand with SymInt implementation #3558

Conversation

miladm commented May 10, 2022 • edited

Action items to unblock:

JackCaoG commented May 10, 2022

Choose a reason for hiding this comment

Choose a reason for hiding this comment

miladm May 10, 2022 • edited

Choose a reason for hiding this comment

miladm commented May 10, 2022

miladm commented May 16, 2022

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

miladm May 16, 2022 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

miladm May 16, 2022 • edited

Choose a reason for hiding this comment

JackCaoG left a comment • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

miladm May 19, 2022 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

miladm May 19, 2022 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

miladm Jul 7, 2022 • edited

Choose a reason for hiding this comment

[POC] Dynamic `expand` with `SymInt` implementation #3558

[POC] Dynamic `expand` with `SymInt` implementation #3558

miladm commented May 10, 2022 •

edited

miladm May 10, 2022 •

edited

miladm May 16, 2022 •

edited

miladm May 16, 2022 •

edited

JackCaoG left a comment •

edited

miladm May 19, 2022 •

edited

miladm May 19, 2022 •

edited

miladm Jul 7, 2022 •

edited