[TIR] not estimating the flops when there is a default estimated flops as attr

[farshidsp]

FlopEstimator is run as part of the Meta Schedule as an estimate for which primfunc to spend time on; e.g. which primfunc to send more trials on. In a case that the for loop extent is an expression, it’s getting stuck on the calculation of the FLOPS.

This PR add a function annotation to the main primfunc and then have FlopEstimator check for this annotation before visiting the function. And if the annotation exists, just use the handcoded flops instead of trying to estimate.

@Lunderberg @csullivan

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @Hzfengsy, @junrushao, @quic-sanirudh, @shingjan _{See #10317 for details}

_{Generated by tvm-bot}

[farshidsp]

@tvm-bot rerun

[Lunderberg]

Looks good to me, just one change requested.

[Lunderberg]

Could we reduce this to a minimal PrimFunc as well? That way, it makes it clear to the reader that the test is validating the behavior of the attribute, and that the if/else/floormod aren't required for the test.

@T.prim_func
def func(A: T.Buffer(1)):
    T.func_attr("estimated_flops": 32)
    for i in T.serial(0, 16):
        A[0] = A[0] + 1

[farshidsp]

Sounds good Eric. Thanks for the suggestion. Fixed in the next commit.

[Lunderberg]

@tvm-bot rerun