[Question][Feature Request] Provides dispatch primitives based on dtypes.

[oraluben]

PyTorch's AT_DISPATCH_FLOATING_TYPES_AND_HALF (and similar ops) are very helpful when implementing kernel for multiple dtypes, flash-infer also implement DISPATCH_DLPACK_*.

It would be great if tvm-ffi could provide those utils.

[tqchen]

@yzh119 @cyx-6 do you think if we can bring FLASHINFER related utils to here?

Something like TVM_FFI_DISPATCH_DTYPE_XXX

[tqchen]

However, would be good to discuss what are desirable canonical name candidates, @oraluben @yzh119 @cyx-6 do u have suggestions?

[oraluben]

While I don't have suggestion on names (TVM_FFI_DISPATCH_XXX sounds good enough), I think the API could be more configurable, and wonder if that makes sense to you.

something like:

TVM_FFI_DISPATCH_DTYPE(input.dtype)
TVM_FFI_DISPATCH_MATCH(DLBfloat16)
TVM_FFI_DISPATCH_MATCH(DLFloat16)
TVM_FFI_DISPATCH_BODY([&] {
    scalar_t input  // scalar_t is template and could be resolved to bfloat16/float16
})

[tqchen]

Not sure how such declarative method can be implemented, usually because it resolves into a switch case, likely we need some builtin categories like FLOATS_AND_HALF etc and pack into a single macro

[tqchen]

so question for @oraluben do u have specific dispatch types in mind, would FLOATS_AND_HALF be sufficient?

[oraluben]

@tqchen yeah, FLOATS_AND_HALF is the only case we met so far. But I think the scenario varies a lot from different users, so I was imaging a declarative and extensible way of doing so.

We workarounded it for our case, so you could take time and see if there's a better API than the torch's. I hope tvm-ffi could maintain the relation between DLDataType and ctypes, and provide composable macros to do so (yes my example is ugly, but I hope it demonstrates the idea).

[yzh119]

do you think if we can bring FLASHINFER related utils to here?

Sure I don't see any problem with it.

Regarding the predefined dispatch macros, pytorch have defined a set of them in https://github.com/pytorch/pytorch/blob/75f798e05b35cbdf1814673e745b9b7817bb29cd/aten/src/ATen/Dispatch.h, but personally I still feel like they are too fragile and now always fit my use case.

[junrushao]

this thread brings me back to the mshadow days :)

[oraluben]

FYI: Just found torch has https://github.com/pytorch/pytorch/blob/9760a633ba9a7d998d88b3bc017c4f8786bd9280/aten/src/ATen/Dispatch_v2.h, if that's helpful.

[junrushao]

It’s a great idea but I feel it a bit under specified. @oraluben could you elaborate which interface you are looking for?

[oraluben]

The core feature I want is the declarative AT_DISPATCH_V2 that supports variable dispatch dtype.