[TensorIR] New schedule primitive set_dtype

[yzh119]

Motivation

Currently, we miss a schedule primitive to change the data type of allocated buffer (e.g. via cache_read/cache_write), and thus we cannot perform type conversion while loading data from global to shared memory.

This PR adds a new schedule primitive set_dtype that follows the interface of set_scope and allows users to customize the allocated buffers' data type.

Example

Before running set_dtype:

@T.prim_func
def before_set_dtype(
    A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")
) -> None:
    B = T.alloc_buffer((128, 128), dtype="float32")

    for i, j in T.grid(128, 128):
        with T.block("B"):
            vi, vj = T.axis.remap("SS", [i, j])
            B[vi, vj] = A[vi, vj] * 2.0
    for i, j in T.grid(128, 128):
        with T.block("C"):
            vi, vj = T.axis.remap("SS", [i, j]
            C[vi, vj] = B[vi, vj] + 1.0

then we perform the set_dtype schedule:

sch = tir.Schedule(before_set_dtype)
sch.set_dtype("B", buffer_index=0, dtype="float16")
print(sch.mod["main"].script())

we get transformed code:

@T.prim_func
def after_set_dtype(
    A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")
) -> None:
    B = T.alloc_buffer((128, 128), dtype="float16")

    for i, j in T.grid(128, 128):
        with T.block("B"):
            vi, vj = T.axis.remap("SS", [i, j])
            B[vi, vj] = T.cast(A[vi, vj] * 2.0, "float16")
    for i, j in T.grid(128, 128):
        with T.block("C"):
            vi, vj = T.axis.remap("SS", [i, j]
            C[vi, vj] = T.cast(B[vi, vj], "float32") + 1.0

where data type conversions are inserted automatically.

Other Usage

Using the combination of cache_read + set_dtype can help us load data from the memory hierarchy while converting data to the desired type.

cc @Hzfengsy @vinx13 @junrushao

[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.

• No users to tag found in teams: tensorir, status: need update _{See #10317 for details}

_{Generated by tvm-bot}

[Hzfengsy]

It's a bit tricky since it would change the behavior of the PrimFunc (aka get different result before and after the schedule)

[quic-sanirudh]

It's a bit tricky since it would change the behavior of the PrimFunc (aka get different result before and after the schedule)

There seems to be some interest in introducing schedule primitives like these that might cause different output. This certainly seems like a useful primitive to have, but since we don't normally want schedule primitives to cause the output to change, primitives like these might be hard to introduce.

Would it make sense to have some discussion on perhaps introducing a different scheduling class that allows these kinds of primitives, or explicitly setting a "danger" flag in the Schedule class before using these primitives. That way, we could still allow primitives like these and make sure that a user knows the risk of using a primitive that could cause the output to change.

[Hzfengsy]

Thanks @quic-sanirudh for the suggestion. I agree that it's a really useful primitive/tool for optimizations. It makes sense if we introduce a danger mode or flag :-)

[yzh119]

Adding a danger flag sounds good to me, any ideas on how to make proper abstractions for danger?

[tqchen]

how about unsafe_set_dtype, with document to users that this operation can change the model behavior.

[junrushao]

@yzh119 any updates

[yzh119]

I have updated my code with new name unsafe_set_dtype.

[tqchen]

oops, @yzh119 please also check the c++ apis @junrushao mentioned :)

[junrushao]

LGTM!

[junrushao]

I don't know why but now it requires committer approval to run github actions for existing contributors. I'm merging this in after the actions got green