Add dynamic support for topk #4644

ymwangg · 2023-02-16T23:52:24Z

This PR adds dynamic shape support for topk.

test/test_dynamic_shapes.py

torch_xla/csrc/aten_xla_type.cpp

torch_xla/csrc/ops/topk.cpp

JackCaoG · 2023-02-17T02:00:55Z

torch_xla/csrc/ops/topk.cpp

+  xla::Shape values_shape =
+      xla::ShapeUtil::MakeShape(input_shape.element_type(), dimensions);
+  xla::Shape indices_shape =
+      xla::ShapeUtil::MakeShape(xla::PrimitiveType::S64, dimensions);


index shape is actully s32, can you use https://github.com/pytorch/xla/blob/master/torch_xla/csrc/tensor_util.cpp#L1254 ? context is SetDimensionShape and getDimensionShpe currently only handle s32, this is where we have slight difference with upstream.

Actually the static version also returns s64 type, because of a s32 to s64 cast here:

xla/torch_xla/csrc/xla_lower_util.cpp

Lines 380 to 383 in 1bbe4da

// aten::topk() wants Long tensors as indices.

return {values, xla::ConvertElementType(

indices, GetDevicePrimitiveType(xla::PrimitiveType::S64,

/*device=*/nullptr))};

Do I also need to disable this cast?

vanbasten23 · 2023-02-17T17:59:51Z

Thanks for the contribution!

test/test_dynamic_shapes.py

vanbasten23 · 2023-02-17T18:11:45Z

torch_xla/csrc/tensor_methods.cpp

      input->CreateFrom(torch::lazy::Value(node, 1), at::ScalarType::Long));
 }

+std::tuple<XLATensorPtr, XLATensorPtr> topk_symint(const XLATensorPtr& input,


quick question, now that we have a symint version of topk, do we still need the original topk?

I basically followed the expand_symint design, which keeps the static expand implementation. I guess the original topk implementation is still needed because the codegen for handling SymInt is slightly different (e.g. SetDimensionSize).

yea, let's keep the static version their.

vanbasten23 · 2023-02-17T18:52:43Z

test/test_dynamic_shapes.py

+    t1[3][1] = 1
+    t2 = torch.nonzero(t1)
+    t3 = torch.zeros([10, 2], device=dev)
+    values, indices = torch.topk(t3, t2.shape[0], dim=0)


I wonder if your change can handle the case where t3 is dynamic, eg t3=torch.nonzero(t1).

I think this feature is already supported by xla through dynamic padder (static topk also works with dynamic t3). I'll add a test case for it later.

ymwangg · 2023-02-21T19:13:20Z

torch_xla/csrc/tensor_methods.cpp

+      largest, sorted, stable);
+  return std::make_tuple(
+      input->CreateFrom(torch::lazy::Value(node, 0)),
+      input->CreateFrom(torch::lazy::Value(node, 1), at::ScalarType::Long));


Some pytorch tests require the static topk returned indices type to be int64, otherwise it throws an error in pytorch here. I saw nonzero returns int32 indices. Why not just do a int64 (pytorch) -> int32 (xla) -> int64 (pytorch) conversion for handling indices?

so what happened here is that underlying XLA type will still be xla::s32, however we mark the logical element type will be at::ScalarType::Long. This will create a problem since underlying s32 can not represent all numbers for long. We run into issues when we try to cast this kind of tensor(long logical type, real s32 type) to other dtypes, because Cast op doesn't know how to handle this kind of cases.

The right thing to do is probably do a manual s64->s32->s64 cast instead of playing the trick of logical_element_type here.

Thanks for the clarification! Do we want to add something like res = torch::lazy::MakeNode<Cast>(res, at::ScalarType::Long) here or use
xla::ConvertElementType(indices, xla::PrimitiveType::S64, /*device=*/nullptr)) in CreateTopK like before?

Let's handle it inside CreateTopK so it is cleaner.

ymwangg · 2023-02-23T20:05:30Z

This PR requires removing torch_pin. Let me ping someone for review.

JackCaoG · 2023-02-23T21:08:52Z

@ymwangg can upstream pr merge without this one? If so we can merge the upstream one first, otherwise we would have to do the two way pining.

ymwangg · 2023-02-23T21:15:55Z

@JackCaoG Yes, the upstream PR does not depend on anything on xla.

Companion PR for pytorch/xla#4644. Pull Request resolved: #95015 Approved by: https://github.com/ezyang

Companion PR for pytorch/xla#4644. Pull Request resolved: pytorch/pytorch#95015 Approved by: https://github.com/ezyang

JackCaoG · 2023-06-28T18:16:28Z

Sorry I have overlooked this pr. @vanbasten23 can you take over this pr and rebase? I think we should merge this pr since it is already here.

Companion PR for pytorch/xla#4644. Pull Request resolved: pytorch#95015 Approved by: https://github.com/ezyang

ymwangg mentioned this pull request Feb 16, 2023

SymIntify topk pytorch/pytorch#95015

Closed

JackCaoG added the dynamism Dynamic Shape Features label Feb 17, 2023

ymwangg force-pushed the topk_symint branch from ed20336 to 8e97967 Compare February 17, 2023 00:11

miladm assigned ymwangg Feb 17, 2023

miladm requested review from miladm and vanbasten23 February 17, 2023 01:40

JackCaoG reviewed Feb 17, 2023

View reviewed changes

test/test_dynamic_shapes.py Show resolved Hide resolved

JackCaoG reviewed Feb 17, 2023

View reviewed changes

torch_xla/csrc/aten_xla_type.cpp Show resolved Hide resolved

JackCaoG reviewed Feb 17, 2023

View reviewed changes

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

JackCaoG reviewed Feb 17, 2023

View reviewed changes

vanbasten23 reviewed Feb 17, 2023

View reviewed changes

test/test_dynamic_shapes.py Outdated Show resolved Hide resolved

vanbasten23 reviewed Feb 17, 2023

View reviewed changes

test/test_dynamic_shapes.py Show resolved Hide resolved

vanbasten23 reviewed Feb 17, 2023

View reviewed changes

ymwangg force-pushed the topk_symint branch from 736dfc1 to def1e04 Compare February 21, 2023 18:52

ymwangg commented Feb 21, 2023

View reviewed changes

ymwangg force-pushed the topk_symint branch from def1e04 to bb6b1a1 Compare February 22, 2023 18:47

vanbasten23 approved these changes Feb 23, 2023

View reviewed changes

pytorchmergebot pushed a commit to pytorch/pytorch that referenced this pull request Feb 24, 2023

SymIntify topk (#95015)

b855b5e

Companion PR for pytorch/xla#4644. Pull Request resolved: #95015 Approved by: https://github.com/ezyang

cyyever pushed a commit to cyyever/pytorch_private that referenced this pull request Feb 25, 2023

SymIntify topk (#95015)

fa9ad40

Companion PR for pytorch/xla#4644. Pull Request resolved: pytorch/pytorch#95015 Approved by: https://github.com/ezyang

cyyever pushed a commit to cyyever/pytorch_private that referenced this pull request Feb 25, 2023

SymIntify topk (#95015)

dd0e7fe

Companion PR for pytorch/xla#4644. Pull Request resolved: pytorch/pytorch#95015 Approved by: https://github.com/ezyang

ymwangg added 4 commits February 28, 2023 04:28

Initial implementation of topk_symint

d180f98

Add .torch_pin

2cb242c

Address CR comments

f4156d4

Remove torch_pin

1f39f88

ymwangg force-pushed the topk_symint branch from bb6b1a1 to 1f39f88 Compare February 28, 2023 04:50

cyyever pushed a commit to cyyever/pytorch_private that referenced this pull request Mar 5, 2023

SymIntify topk (#95015)

0b08472

Companion PR for pytorch/xla#4644. Pull Request resolved: pytorch/pytorch#95015 Approved by: https://github.com/ezyang

jhavukainen pushed a commit to kulinseth/pytorch that referenced this pull request Mar 15, 2024

SymIntify topk (pytorch#95015)

56d5ba4

Companion PR for pytorch/xla#4644. Pull Request resolved: pytorch#95015 Approved by: https://github.com/ezyang

	// aten::topk() wants Long tensors as indices.
	return {values, xla::ConvertElementType(
	indices, GetDevicePrimitiveType(xla::PrimitiveType::S64,
	/device=/nullptr))};

Add dynamic support for topk #4644

Are you sure you want to change the base?

Add dynamic support for topk #4644

Uh oh!

Conversation

ymwangg commented Feb 16, 2023

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!

vanbasten23 commented Feb 17, 2023

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!

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!

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!

ymwangg commented Feb 23, 2023

Uh oh!

JackCaoG commented Feb 23, 2023

Uh oh!

ymwangg commented Feb 23, 2023

Uh oh!

JackCaoG commented Jun 28, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants