[POC] Dynamic `expand` with `SymInt` implementation by miladm · Pull Request #3558 · pytorch/xla

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:

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

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

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

+  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

+
+    /* 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?

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

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

+  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

+  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.

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

+                        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

  - 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

-  -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

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

+
+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

+  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

+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..

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

+    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 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

Comment thread torch_xla/csrc/aten_xla_type.cpp Outdated

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

Comment thread torch_xla/csrc/ops/ops.cpp Outdated

JackCaoG reviewed May 10, 2022

View reviewed changes

Comment thread torch_xla/csrc/ops/ops.cpp Outdated

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

Comment thread torch_xla/csrc/tensor_methods.cpp Outdated

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 20 commits July 4, 2022 06:32

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

ysiraichi added DO_NOT_MERGE Not for merging. and removed DO_NOT_MERGE_YET labels Mar 5, 2025


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

		SizeAdd::SizeAdd(XlaValue a, XlaValue b)

Conversation

miladm commented May 10, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Action items to unblock:

Uh oh!

Uh oh!

JackCaoG commented May 10, 2022

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

miladm May 10, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

miladm commented May 10, 2022

Uh oh!

Uh oh!

Uh oh!

miladm commented May 16, 2022

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

miladm May 16, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

miladm May 16, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JackCaoG left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

miladm May 19, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

miladm May 19, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

miladm Jul 7, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

miladm commented May 10, 2022 •

edited

Loading

miladm May 10, 2022 •

edited

Loading

miladm May 16, 2022 •

edited

Loading

miladm May 16, 2022 •

edited

Loading

JackCaoG left a comment •

edited

Loading

miladm May 19, 2022 •

edited

Loading

miladm May 19, 2022 •

edited

Loading

miladm Jul 7, 2022 •

edited

Loading