Enable passing down dynamic dimensions from torch to XLA #5790

lsy323 · 2023-11-11T03:34:56Z

This is based on #5778. This PR focuses on supporting unbounded dynamism in LTC infra. The e2e would work when the unbounded dynamism support on XLA side is merged, and required op lowering logic is updated in torch_xla.

Add support in LTC to pass down the unbounded dynamism info down to HLO lowering.
Guard op lowering logic change for unbounded dynamism support by EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM=1
Fixing a bug in setup.py

Test:
Tested with @sdasgup3's XLA change to propagate unbounded dynamism branch. To make the patch works with torch_xla HEAD, the branch needs to be rebased to the pinned XLA version. Rebased branch

The unbounded dynamism can be propagated in the following example:

import torch
import torch_xla
import torch_xla.core.xla_model as xm

device = xm.xla_device()

a = torch.tensor([[1,2],[2,4]], device=device)
torch_xla._XLAC._xla_mark_dynamic(a, 0)
b = torch.tensor([[1,2],[2,4]], device=device)
torch_xla._XLAC._xla_mark_dynamic(b, 0)
c = a * b
hlo_content = torch_xla._XLAC._get_xla_tensors_hlo([c])
print(hlo_content)

which gives

loModule IrToHlo.5, entry_computation_layout={(s64[?,2]{1,0}, s64[?,2]{1,0})->(s64[?,2]{1,0})}

ENTRY %IrToHlo.5 (p0.1: s64[?,2], p1.2: s64[?,2]) -> (s64[?,2]) {
  %p1.2 = s64[?,2]{1,0} parameter(1)
  %p0.1 = s64[?,2]{1,0} parameter(0)
  %multiply.3 = s64[?,2]{1,0} multiply(s64[?,2]{1,0} %p1.2, s64[?,2]{1,0} %p0.1)
  ROOT %tuple.4 = (s64[?,2]{1,0}) tuple(s64[?,2]{1,0} %multiply.3)
}

qihqi

Please change the title of PR to something more descriptive: something like "Enable passing down dynamic dimensions from torch to XLA" etc.

Also all C++ use of macro change that to env var like the python one. (so that no need to recompile to enable the feature)

torch_xla/csrc/elementwise.cpp

JackCaoG · 2023-11-13T19:02:52Z

torch_xla/csrc/helpers.cpp

@@ -322,6 +329,116 @@ xla::XlaOp XlaHelpers::DynamicReshapeAs(xla::XlaOp input,
             : xla::Reshape(input, shape.dimensions());
 }

+#if EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM


not a fan of macro, if you want to guard these functions you can add a check like XLA_CHECK(getenv(EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM))

sg, will update

torch_xla/csrc/helpers.cpp

JackCaoG · 2023-11-13T19:20:08Z

torch_xla/csrc/helpers.cpp

+        get_dim_ops.push_back(xla::GetDimensionSize(aux_op, i));
+
+        auto s = ShapeHelper::ShapeOfXlaOp(get_dim_ops.back());
+        std::cout << "implicitB shape: " << xla::ShapeUtil::HumanString(s)


torch_xla/csrc/helpers.cpp

torch_xla/csrc/ir.h

torch_xla/csrc/lowering_context.cpp

JackCaoG · 2023-11-13T19:36:04Z

torch_xla/csrc/lowering_context.cpp

+    xla::internal::XlaBuilderFriend builder_friend;
+    auto* inst = builder_friend.GetInstruction(result_ops[0]);
+    auto* mutable_dynamic =
+        inst->mutable_shape()->mutable_is_dynamic_dimension();
+    if (mutable_dynamic->empty()) {
+      for (int i = 0; i < inst->dimensions_size(); i++) {
+        mutable_dynamic->Add(false);
+      }
+    }
+    auto* mutable_dims = inst->mutable_shape()->mutable_dimensions();
+    for (const auto dim : casted->dynamic_dims()) {
+      mutable_dynamic->Set(dim, true);
+      mutable_dims->Set(dim, kUnboundedSize);
+    }


do we always need to run these or we should guard this in a conditional?

Let's guard it by checking if the XlaNode has dynamic dim

…d dynamism code path

torch_xla/csrc/helpers.cpp

miladm · 2023-11-14T20:05:19Z

torch_xla/csrc/elementwise.cpp

-                             0, input_shape.element_type(), input.builder()));
+  xla::XlaOp scalar = XlaHelpers::ScalarValue<float>(
+      0, input_shape.element_type(), input.builder());
+  if (experimental_unbounded_dynamism) {


Can this path execute in the non export path?
Wdyt we limit this experimental condition to torch.export path?

Per the discussion below, seems there is no better solution than using a env variable to enable the unbounded dynamism lowering. But we can put more thought on this one.

SG. This issue doesn't block this PR. Let's prepare a proposal in parallel.

miladm · 2023-11-14T20:08:21Z

torch_xla/csrc/helpers.cpp

+                     [](int64_t size) { return size == kUnboundedSize; });
+}
+
+xla::XlaOp XlaHelpers::DynamicUnboundedReshape(


Can you please prepare an opset issue to track supporting this mode of dynamism?

Similar work example: #5764

Sure, I think we have a list internally

Does it work to maintain the list on GH?

miladm · 2023-11-14T20:11:31Z

torch_xla/csrc/reduction.cpp

@@ -31,6 +31,9 @@ struct SummationResult {
  xla::XlaOp result;
 };

+static const bool experimental_unbounded_dynamism =
+    runtime::sys_util::GetEnvBool("EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM", false);


env variables causes poor user experience. What's the plan to clean up a better solution (potentially via the upstream torch API level)?

cc @JackCaoG @qihqi @lsy323

Also, for adding env variables, please keep this file up to date.

https://github.com/pytorch/xla/blob/master/configuration.yaml

Current thought is to use env variable to limit the code path to "Export Only", as you've pointed out in many places. If there are better mechanisms to accomplish that we can also use something different.

miladm

Thanks @lsy323 - I left a few comments.

torch_xla/csrc/helpers.cpp

torch_xla/csrc/ir.cpp

torch_xla/csrc/ir.h

torch_xla/stablehlo.py

JackCaoG

mostly lgtm, minor nits

miladm

LGTM. Thanks.

* port sandeep unbounded dynamism change * Enable unbounded dynamism using env var, add more guards for unbounded dynamism code path --------- Co-authored-by: Siyuan Liu <[email protected]>

Siyuan Liu and others added 4 commits November 11, 2023 03:26

port sandeep unbounded dynamism change

56f698d

format

24f42c4

fix linter

4acf02d

remove set tag

495f844

lsy323 marked this pull request as ready for review November 13, 2023 17:52

lsy323 requested review from qihqi and JackCaoG November 13, 2023 17:53

qihqi approved these changes Nov 13, 2023

View reviewed changes

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

lsy323 changed the title ~~Mirror of #5778 Make code compilable~~ Enable passing down dynamic dimensions from torch to XLA Nov 13, 2023