Function check_region_bound

tests/python/arith/test_arith_intset.py:140–208

Helper to check region bound estimation. Parameters ---------- expect_region: dict The keys are of form (begin, end) or PrimExpr as a single point. The values are expected estimated region or region dict on different bindings. var_dom: dict Map var to iterat

(expect_region, var_dom, mode, predicate=None)

def check_region_bound(expect_region, var_dom, mode, predicate=None):
    """Helper to check region bound estimation.

    Parameters
    ----------
    expect_region: dict
        The keys are of form (begin, end) or PrimExpr as a single point. The values are
        expected estimated region or region dict on different bindings.

    var_dom: dict
        Map var to iteration domain range.

    mode: str
        Specify "lowerbound", "upperbound" or else use strict bound estimation.

    predicate: PrimExpr
        Extra predicate, defaults to True.
    """
    if predicate is None:
        predicate = tvm.tirx.IntImm("bool", 1)
    region = []
    expect = []
    for k, v in expect_region.items():
        if not isinstance(k, tuple | list):
            k = (k, k + 1)
        region.append(tvm.ir.Range.from_min_extent(k[0], Analyzer().simplify(k[1] - k[0])))
        expect.append(v)
    if mode == "lowerbound":
        result = tvm.arith.estimate_region_lower_bound(
            region=region, var_dom=var_dom, predicate=predicate
        )
    elif mode == "upperbound":
        result = tvm.arith.estimate_region_upper_bound(
            region=region, var_dom=var_dom, predicate=predicate
        )
    else:
        result = tvm.arith.estimate_region_strict_bound(
            region=region, var_dom=var_dom, predicate=predicate
        )
    if result is None:
        assert all([_ is None for _ in expect])
        return
    assert len(result) == len(expect)
    for intset, expect_desc in zip(result, expect):
        if isinstance(expect_desc, dict):
            # check range on different free var bindings
            for binding in expect_desc:
                analyzer = Analyzer()
                for k, v in binding:
                    analyzer.bind(k, v)
                expect_begin, expect_end = expect_desc[binding]
                result_begin = analyzer.simplify(intset.min_value, 3)
                result_end = analyzer.simplify(intset.max_value + 1, 3)
                assert analyzer.can_prove_equal(result_begin - expect_begin, 0), (
                    f"{result_begin} vs {expect_begin}"
                )
                assert analyzer.can_prove_equal(result_end - expect_end, 0), (
                    f"{result_end} vs {expect_end}"