Add support for `select` op #2179

zasdfgbnm · 2022-11-11T07:56:02Z

Select op takes a slice of a tensor at a given dimension. For example, for a 3D tensor, t.select(1, 123) means taking a slice along dim 1 at index 123, which is identical to t[:, 123, :].

The reason why I write this PR is:

select is a commonly used op in PyTorch, it is always good to have better op coverage.
It is similar to index_select but easier, so this PR could serve as a reference for the index_select, and it is also a good place to discuss indexing about index_select here.

The current support of select is not complete, it has restrictions that require the input tensor to be a fusion input. And I don't think in registry.cpp, this PR is detecting all unsupported cases and rejecting them. But for now, I believe providing this level of support is sufficient.

torch/csrc/jit/codegen/cuda/lower_index.cpp

torch/csrc/jit/codegen/cuda/ir_utils.cpp

torch/csrc/jit/codegen/cuda/index_compute.cpp

torch/csrc/jit/codegen/cuda/scheduler/registry.cpp

naoyam

LGTM. Thanks for adding the support of select.

torch/csrc/jit/codegen/cuda/arith.h

torch/csrc/jit/codegen/cuda/mutator.cpp

torch/csrc/jit/codegen/cuda/scheduler/utils.cpp

naoyam · 2022-11-11T18:54:33Z

torch/csrc/jit/codegen/cuda/index_compute.cpp

@@ -1615,7 +1612,10 @@ std::vector<Val*> Index::getGlobalProducerStridedIndices(
        loops,
        root_ind);

-    root_ind = getProducerIndexWithHalo(producer_tv, i, root_ind, consumer_tv);
+    if (root_dom[i] != selected_id) {


I think we should just assert this for now. getProducerIndexWithHalo should just return root_ind if the root domain is not extended with halo nor accessed through a shift op.

There is an assert inside getProducerHaloOffset:

auto it = p2c.find(producer_id); // p2c should always have a mapping for producer_id. The only case // where no mapping exists for a producer axis is when it is a // reduction axis. Since this function is only used for indexing // producer tensors, where reduction axes are skipped, producer_id // should never be a reduction axis. TORCH_INTERNAL_ASSERT(it != p2c.end());

which fails for the select id. Do we want to remove this assert and return 0 if not mapped?

Yeah, I think we should remove the assert, but I'd feel more assured with some validation that the producer ID is a selected ID.

I moved the check inside getProducerIndexWithHalo, something like

if (override) { offset = 0; } else { assert }

naoyam

LGTM. Added some minor comments.

torch/csrc/jit/codegen/cuda/lower_index.h

naoyam · 2022-11-14T06:31:42Z

torch/csrc/jit/codegen/cuda/index_compute.cpp

@@ -1591,18 +1575,27 @@ std::vector<Val*> Index::getGlobalProducerStridedIndices(
      continue;
    }

-    TORCH_INTERNAL_ASSERT(
+    Val* root_ind = nullptr;
+    auto override_it = override_index.find(root_dom[i]);


I was thinking about passing the optional map to getTensorIndexFromIdGraph to provide an initial ID-to-index map. That would be more consistent if we would want to allow the same initial map in consumer indexing.

That said, I think this is good enough for now given that the whole indexing code would be redesigned.

Pinging @csarofeen

naoyam · 2022-11-14T06:45:51Z

torch/csrc/jit/codegen/cuda/index_compute.cpp

@@ -1615,7 +1612,10 @@ std::vector<Val*> Index::getGlobalProducerStridedIndices(
        loops,
        root_ind);

-    root_ind = getProducerIndexWithHalo(producer_tv, i, root_ind, consumer_tv);
+    if (root_dom[i] != selected_id) {


Yeah, I think we should remove the assert, but I'd feel more assured with some validation that the producer ID is a selected ID.

torch/csrc/jit/codegen/cuda/index_compute.cpp

csarofeen · 2022-11-16T00:24:53Z

Is select an alias operation?

zasdfgbnm · 2022-11-16T00:29:57Z

Is select an alias operation?

It is, it returns a view of the input tensor.

naoyam · 2022-11-16T02:07:55Z

Is select an alias operation?

It is, it returns a view of the input tensor.

In nvFuser, even a view isn't really an alias as it creates a new (transformed) copy. I don't think we have a true alias in nvFuser. Is it different in PyTorch? Is a view a true "view" to another tensor?

zasdfgbnm · 2022-11-16T02:39:16Z

Is select an alias operation?

It is, it returns a view of the input tensor.

In nvFuser, even a view isn't really an alias as it creates a new (transformed) copy. I don't think we have a true alias in nvFuser. Is it different in PyTorch? Is a view a true "view" to another tensor?

Yes, in eager mode, view is a true view. If the stride is not compatible with the viewed shape, then view will throw an error. For this case, user can use reshape to make a copy of the tensor with compatible stride.

naoyam · 2022-11-16T04:16:47Z

Is select an alias operation?

It is, it returns a view of the input tensor.

In nvFuser, even a view isn't really an alias as it creates a new (transformed) copy. I don't think we have a true alias in nvFuser. Is it different in PyTorch? Is a view a true "view" to another tensor?

Yes, in eager mode, view is a true view. If the stride is not compatible with the viewed shape, then view will throw an error. For this case, user can use reshape to make a copy of the tensor with compatible stride.

Hmm, so I'm assuming we somehow handle the semantic difference, right? I wonder what would happen when writing to view tensors.

zasdfgbnm · 2022-11-16T05:40:01Z

Is select an alias operation?

It is, it returns a view of the input tensor.

In nvFuser, even a view isn't really an alias as it creates a new (transformed) copy. I don't think we have a true alias in nvFuser. Is it different in PyTorch? Is a view a true "view" to another tensor?

Yes, in eager mode, view is a true view. If the stride is not compatible with the viewed shape, then view will throw an error. For this case, user can use reshape to make a copy of the tensor with compatible stride.

Hmm, so I'm assuming we somehow handle the semantic difference, right? I wonder what would happen when writing to view tensors.

I don't know. I remember seeing something called view_copy, not sure how it works though. Probably TorchScript might convert the IR to SSA before feeding to nvfuser? @jjsjann123

csarofeen · 2022-11-16T19:05:57Z

Yeah, this was why I asked. We would need to figure out how to support the alias analysis of the op to pipe it through.

jjsjann123 · 2022-11-16T19:17:29Z

sounds like we needed to have an alias analysis to ensure that changing the select op to select_copy preserves correct semantics. We did that for other view like ops.

zasdfgbnm added 2 commits November 10, 2022 23:47

Add support for select op

b8a8c29

save

caff3ac

zasdfgbnm marked this pull request as draft November 11, 2022 07:56

zasdfgbnm added 2 commits November 11, 2022 00:15

reduction working

ee5b0e8

persistent working

2439478

zasdfgbnm commented Nov 11, 2022

View reviewed changes

torch/csrc/jit/codegen/cuda/lower_index.cpp Outdated Show resolved Hide resolved

zasdfgbnm added 2 commits November 11, 2022 00:38

remove dead code

041dbf1

format

11a56e6

zasdfgbnm commented Nov 11, 2022

View reviewed changes

torch/csrc/jit/codegen/cuda/ir_utils.cpp Show resolved Hide resolved

zasdfgbnm marked this pull request as ready for review November 11, 2022 08:58

zasdfgbnm requested a review from naoyam November 11, 2022 08:58

zasdfgbnm commented Nov 11, 2022

View reviewed changes

torch/csrc/jit/codegen/cuda/index_compute.cpp Show resolved Hide resolved

zasdfgbnm commented Nov 11, 2022

View reviewed changes

torch/csrc/jit/codegen/cuda/scheduler/registry.cpp Show resolved Hide resolved

naoyam approved these changes Nov 11, 2022

View reviewed changes

refactor

87fee1d

zasdfgbnm requested a review from naoyam November 11, 2022 20:46

zasdfgbnm added 2 commits November 11, 2022 14:32

fix shift

5a54098

cacheInputs

02b2425

naoyam approved these changes Nov 14, 2022

View reviewed changes

naoyam mentioned this pull request Nov 15, 2022

Add support for index_select OP (backend implementation) #2183

Merged

zasdfgbnm added 2 commits November 14, 2022 22:15

Merge branch 'devel' of github.com:csarofeen/pytorch into select

3ae76f5

review comments

7ee9b0b

zasdfgbnm merged commit ca1387a into devel Nov 15, 2022

zasdfgbnm deleted the select branch November 15, 2022 07:42

Add support for select op #2179

Add support for select op #2179

Uh oh!

Conversation

zasdfgbnm commented Nov 11, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

naoyam left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

naoyam Nov 11, 2022

Choose a reason for hiding this comment

Uh oh!

zasdfgbnm Nov 11, 2022

Choose a reason for hiding this comment

Uh oh!

naoyam Nov 14, 2022

Choose a reason for hiding this comment

Uh oh!

zasdfgbnm Nov 15, 2022

Choose a reason for hiding this comment

Uh oh!

naoyam left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

naoyam Nov 14, 2022

Choose a reason for hiding this comment

Uh oh!

csarofeen Nov 16, 2022

Choose a reason for hiding this comment

Uh oh!

naoyam Nov 14, 2022

Choose a reason for hiding this comment

Uh oh!

Uh oh!

csarofeen commented Nov 16, 2022

Uh oh!

zasdfgbnm commented Nov 16, 2022

Uh oh!

naoyam commented Nov 16, 2022

Uh oh!

zasdfgbnm commented Nov 16, 2022

Uh oh!

naoyam commented Nov 16, 2022

Uh oh!

zasdfgbnm commented Nov 16, 2022

Uh oh!

csarofeen commented Nov 16, 2022

Uh oh!

jjsjann123 commented Nov 16, 2022

Uh oh!

Uh oh!

Add support for `select` op #2179

Add support for `select` op #2179

zasdfgbnm commented Nov 11, 2022 •

edited

Loading