[Metaschedule] Add test case for multi-anchor subgraph

[masahi]

This adds a demonstration of extracting, scheduling, and e2e-compiling relay subgraphs with multiple anchor ops. Since task extraction is not associated with TE scheduling anymore, extracting a subgraph with multiple anchor TE compute just works.

The test case manually creates a simple fused mod with two relay.dense. But in the future, an effort like #9628 should make it easier to construct multi-anchor subgraphs.

The extracted TensorIR block corresponding to two TE dense compute looks like this:

@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(placeholder: T.Buffer[(128, 128), "float32"], placeholder_1: T.Buffer[(128, 128), "float32"], placeholder_2: T.Buffer[(128, 128), "float32"], T_matmul_NT: T.Buffer[(128, 128), "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        T_matmul_NT_1 = T.alloc_buffer([128, 128], dtype="float32")
        for i0, i1, i2 in T.grid(128, 128, 128):
            with T.block("T_matmul_NT"):
                i, j, k = T.axis.remap("SSR", [i0, i1, i2])
                T.reads(placeholder[i, k], placeholder_1[j, k])
                T.writes(T_matmul_NT_1[i, j])
                T.block_attr({"layout_free_placeholders":[placeholder_1]})
                with T.init():
                    T_matmul_NT_1[i, j] = T.float32(0)
                T_matmul_NT_1[i, j] = T_matmul_NT_1[i, j] + placeholder[i, k] * placeholder_1[j, k]
        for i0, i1, i2 in T.grid(128, 128, 128):
            with T.block("T_matmul_NT_1"):
                i, j, k = T.axis.remap("SSR", [i0, i1, i2])
                T.reads(T_matmul_NT_1[i, k], placeholder_2[j, k])
                T.writes(T_matmul_NT[i, j])
                T.block_attr({"layout_free_placeholders":[placeholder_2]})
                with T.init():
                    T_matmul_NT[i, j] = T.float32(0)
                T_matmul_NT[i, j] = T_matmul_NT[i, j] + T_matmul_NT_1[i, k] * placeholder_2[j, k]
    

@junrushao1994 @csullivan @comaniac @mbs-octoml @mikepapadim

[comaniac]

LGTM

[masahi]

@junrushao1994 @zxybazh

I keep writing this database boilerplate for manual scheduling. I'm thinking about a clean API so that users don't have to go through the explicit task extraction -> database creation steps. Right now it looks like

    relay_mod = tvm.IRModule.from_expr(...)
    target = "llvm"    
    params = {"weight1": weight1_np, "weight2": weight2_np}

    def schedule_fn(task, sch):
        if "nn_dense_nn_dense" in task.task_name:
            schedule_dense_dense(sch)
            return True
        return False

    database = apply_manual_schedules(relay_mod, target, params, schedule_fn)

    with ApplyHistoryBest(database):
           ...

If this looks ok, I can PR it after this one.

[junrushao]

There is a DummyDatabase class in python/tvm/meta_schedule/testing/utils.py where we don't need to create json files for intermediate results, and I was wondering if we could further reduce boilerplate by enhancing that class. What do you think?

[zxybazh]

Yeah using the Dummy classes can be very helpful in tuning and it's acutally a good idea to provide new use interface as you mentioned, a schedule function to work on on the relay level. Let me know when you got the PR ready : )

[tmoreau89]

LGTM!

[junrushao]

Niiiiiice work! Thanks @masahi!