Skip to content

Fuse reshape + gather op into single gather op #8329

@ddilbazTT

Description

@ddilbazTT

Investigate whether it's possible to fuse and how it should be done when there exists a reshape + gather, like the following:

    %4 = stablehlo.reshape %3 : (tensor<1836732xi64>) -> tensor<1836732x1xi64> loc(#loc)
    %5 = "stablehlo.gather"(%arg1, %4) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, slice_sizes = array<i64: 1, 80>}> : (tensor<1993728x80xf32>, tensor<1836732x1xi64>) -> tensor<1836732x80xf32> loc(#loc)

Ideally, the new gather op would use the original shape: tensor<1836732xi64>

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type
No fields configured for issues without a type.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions