xla icon indicating copy to clipboard operation
xla copied to clipboard

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

Open miladm opened this issue 3 years ago • 3 comments

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

Action items to unblock:

  • [x] Lower expand with SymInt input parameter as a POC. This won't compile until expand with SymInt is available in PyTorch
  • [x] Update upstream PyTorch when expand.SymInt op (with c10::SymIntArrayRef input signature is available) is ready
  • [x] Integrate dynamic_ir as a subclass of XLANode
  • [x] Implement SizeNode lowering using xla::GetDimensionSize API
  • [x] Verify/implement JIT SSA shape support for expand.SymInt in PyTorch
  • [x] Verify/implement shape inference support for expand.SymInt in PyTorch
    • PR: https://github.com/pytorch/pytorch/pull/77830
    • Status: Blocked
  • [x] Reimplement DimensionNode::isDynamic() after PyTorch API support becomes available
    • https://github.com/pytorch/pytorch/issues/77909
    • Support the same feature in PyTorch/XLA (this PR)
  • [x] Implement SymInt-related helper functions to improve reuse and modularity.
  • [x] Move DimensionNode class to LTC core and enable multiple inheritance for Size nodes
    • https://github.com/pytorch/pytorch/pull/78088
  • [x] [Testing] Add a C++ unit test for expand.SymInt in PyTorch/XLA
    • https://github.com/pytorch/xla/issues/3589
    • Temporarily unblocked
  • [x] [Testing] implement support for torch.nonzero in PyTorch
    • PR: https://github.com/pytorch/pytorch/pull/77572
    • Status: Merged
  • [x] [Testing] implement support for torch.nonzero in PyTorch/XLA
  • [x] is_symbolic API issue
    • https://github.com/pytorch/xla/issues/3680

miladm avatar May 10 '22 04:05 miladm

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

JackCaoG avatar May 10 '22 17:05 JackCaoG

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 avatar May 10 '22 18:05 miladm

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 avatar May 16 '22 05:05 miladm