Skip to content

[Bug] bijective_layout.forward_shape losses information of concrete dim #13480

@lixiaoquan

Description

@lixiaoquan

Consider this case:

def test_bilayout_with_any():
    bilayout = tvm.tir.bijective_layout("NCHW", "NHWC")
    assert isinstance(bilayout, tvm.tir.BijectiveLayout)
    dst_shape = bilayout.forward_shape((Any(), 32, 7, Any()))
    print(dst_shape)
    # output is [?, 7, ?, ?]
    src_shape = bilayout.backward_shape(dst_shape)
    print(src_shape)
    # output is [?, ?, ?, ?]

Expected behavior

forward_shape's output should be [?, 7, ?, 32]
backward_shape's output should be [?, 32, 7, ?]

Actual behavior

forward_shape's output should be [?, 7, ?, ?]
backward_shape's output should be [?, ?, ?, ?]

As you can see, all dims become Any().

This is caused by https://github.com/apache/tvm/blob/main/src/tir/ir/data_layout.cc#L372

If original dim is Any(), dim in transformed shape will be Any(). I'm not sure this is the expected behavior. Is there any reason behind this code?

It is introduced by #4179 @kevinthesun Could you please comment?

Environment

Latest TVM

Steps to reproduce

Please check the test above

Triage

Please refer to the list of label tags here to find the relevant tags and add them below in a bullet format (example below).

  • needs-triage

Metadata

Metadata

Assignees

No one assigned

    Labels

    needs-triagePRs or issues that need to be investigated by maintainers to find the right assignees to address ittype: bug

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions