[Bug] Metaschedule flop estimation cannot handle non-integer loop dimensions

[tkonolige]

Metaschedule does not handle non-integer loop extents. This problem is specific to metaschedule as auto_scheduler handles non-integer loop extents. This failure happens with extents that can be simplified to a constant (4+4) as well as with extents that are variable in size (i+4). I think the solution would be to use an arith::Analyzer as auto_schedule does. Or to just use auto_scheduler's flop estimation.

from tvm.ir import IRModule
from tvm.meta_schedule.testing.te_workload import create_te_workload
from tvm.script import tir as T
from tvm.tir.analysis import estimate_tir_flops


@T.prim_func
def flops_with_nonint_extent(a: T.Buffer[16, "float32"]):
    for i in range(4 + 4):
        a[i] = a[i]


def test_flops_with_nonint_extent():
    estimate_tir_flops(IRModule({"main": flops_with_nonint_extent}))


@T.prim_func
def flops_with_variable_extent(a: T.Buffer[16, "float32"]):
    for i in range(4 + 4):
        for j in range(i + 8):
            a[j] = a[i]


def test_flops_with_variable_extent():
    estimate_tir_flops(IRModule({"main": flops_with_variable_extent}))


if __name__ == "__main__":
    test_flops_with_nonint_extent()
    test_flops_with_variable_extent()

@junrushao1994 @jwfromm