-
Notifications
You must be signed in to change notification settings - Fork 480
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[MLIR][TORCH] Add E2E support for aten.index_select op
This commit adds lowering of `aten.index_select` op. Signed-Off By: Vivek Khandelwal <[email protected]>
- Loading branch information
1 parent
0a0a1b4
commit 8130354
Showing
3 changed files
with
224 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,145 @@ | ||
# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. | ||
# See https://llvm.org/LICENSE.txt for license information. | ||
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | ||
# Also available under a BSD-style license. See LICENSE. | ||
|
||
import torch | ||
|
||
from torch_mlir_e2e_test.torchscript.framework import TestUtils | ||
from torch_mlir_e2e_test.torchscript.registry import register_test_case | ||
from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export | ||
|
||
# ============================================================================== | ||
|
||
|
||
class IndexSelectSingleIdxModule(torch.nn.Module): | ||
def __init__(self): | ||
super().__init__() | ||
|
||
@export | ||
@annotate_args([ | ||
None, | ||
([4, 5, 6], torch.float32, True), | ||
([1], torch.int64, True), | ||
]) | ||
|
||
def forward(self, input, indices): | ||
return torch.index_select(input, 1, indices) | ||
|
||
@register_test_case(module_factory=lambda: IndexSelectSingleIdxModule()) | ||
def IndexSelectSingleIdxModule_basic(module, tu: TestUtils): | ||
module.forward(torch.randn(4, 5, 6), torch.tensor([2])) | ||
|
||
|
||
class IndexSelectTwoIdxModule(torch.nn.Module): | ||
def __init__(self): | ||
super().__init__() | ||
|
||
@export | ||
@annotate_args([ | ||
None, | ||
([4, 5, 6], torch.float32, True), | ||
([2], torch.int64, True), | ||
]) | ||
|
||
def forward(self, input, indices): | ||
return torch.index_select(input, 2, indices) | ||
|
||
@register_test_case(module_factory=lambda: IndexSelectTwoIdxModule()) | ||
def IndexSelectTwoIdxModule_basic(module, tu: TestUtils): | ||
module.forward(torch.randn(4, 5, 6), torch.tensor([2, 4])) | ||
|
||
|
||
class IndexSelectWholeDimensionModule(torch.nn.Module): | ||
def __init__(self): | ||
super().__init__() | ||
|
||
@export | ||
@annotate_args([ | ||
None, | ||
([4, 5, 6], torch.float32, True), | ||
([4], torch.int64, True), | ||
]) | ||
|
||
def forward(self, input, indices): | ||
return torch.index_select(input, 0, indices) | ||
|
||
@register_test_case(module_factory=lambda: IndexSelectWholeDimensionModule()) | ||
def IndexSelectWholeDimensionModule_basic(module, tu: TestUtils): | ||
module.forward(torch.randn(4, 5, 6), torch.tensor([0, 1, 2, 3])) | ||
|
||
|
||
class IndexSelectWholeTensorModule(torch.nn.Module): | ||
def __init__(self): | ||
super().__init__() | ||
|
||
@export | ||
@annotate_args([ | ||
None, | ||
([3], torch.float32, True), | ||
([3], torch.int64, True), | ||
]) | ||
|
||
def forward(self, input, indices): | ||
return torch.index_select(input, 0, indices) | ||
|
||
@register_test_case(module_factory=lambda: IndexSelectWholeTensorModule()) | ||
def IndexSelectWholeTensorModule_basic(module, tu: TestUtils): | ||
module.forward(torch.randn(3), torch.tensor([0, 1, 2])) | ||
|
||
|
||
class IndexSelectDynamicModule(torch.nn.Module): | ||
def __init__(self): | ||
super().__init__() | ||
|
||
@export | ||
@annotate_args([ | ||
None, | ||
([-1, -1, -1], torch.float32, True), | ||
([-1], torch.int64, True), | ||
]) | ||
|
||
def forward(self, input, indices): | ||
return torch.index_select(input, 2, indices) | ||
|
||
@register_test_case(module_factory=lambda: IndexSelectDynamicModule()) | ||
def IndexSelectDynamicModulebasic(module, tu: TestUtils): | ||
module.forward(torch.randn(4, 5, 6), torch.tensor([0, 4])) | ||
|
||
|
||
class IndexSelectDynamicInputSizeModule(torch.nn.Module): | ||
def __init__(self): | ||
super().__init__() | ||
|
||
@export | ||
@annotate_args([ | ||
None, | ||
([-1, -1, -1], torch.float32, True), | ||
([2], torch.int64, True), | ||
]) | ||
|
||
def forward(self, input, indices): | ||
return torch.index_select(input, 2, indices) | ||
|
||
@register_test_case(module_factory=lambda: IndexSelectDynamicInputSizeModule()) | ||
def IndexSelectDynamicInputSizeModule_basic(module, tu: TestUtils): | ||
module.forward(torch.randn(4, 5, 6), torch.tensor([0, 2])) | ||
|
||
|
||
class IndexSelectDynamicIndexSizeModule(torch.nn.Module): | ||
def __init__(self): | ||
super().__init__() | ||
|
||
@export | ||
@annotate_args([ | ||
None, | ||
([4, 5, 6], torch.float32, True), | ||
([-1], torch.int64, True), | ||
]) | ||
|
||
def forward(self, input, indices): | ||
return torch.index_select(input, 1, indices) | ||
|
||
@register_test_case(module_factory=lambda: IndexSelectDynamicIndexSizeModule()) | ||
def IndexSelectDynamicIndexSizeModule_basic(module, tu: TestUtils): | ||
module.forward(torch.randn(4, 5, 6), torch.tensor([1, 2])) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters