As a follow-up to my question on discord, I started working on a (rather minimal) approach to implement a .to_launch_arg() method for structs that are annotated with the CubeLaunch macro. The idea is that one can use this similarly to the .to_arg() function for Tensors, making it easier to use a custom struct as an argument when launching CubeCL kernels.
I was not entirely sure where to place all of my code, here is a summary of my approach:
- in
cubecl-core, I added a new trait AsLaunchArgument defining the signature of the as_launch_arg() function
- in
cubecl-macros, I added a new macro called CubeLaunchArg that implements AsLaunchArgument for a struct
- my implementation checks whether a field is a scalar (wrote a function
TypeField::is_scalar(&Type) for this); if so, wrap the field in ScalarArg, if not, hope that the field implements AsLaunchArg.
This is very rudimentary, and is just enough to make my use case work. I would be happy to extend it further, but I think it makes sense to check out whether my approach makes sense to you (as in the CubeCL maintainers & devs).
Thanks!
Marcel
As a follow-up to my question on discord, I started working on a (rather minimal) approach to implement a
.to_launch_arg()method for structs that are annotated with theCubeLaunchmacro. The idea is that one can use this similarly to the.to_arg()function for Tensors, making it easier to use a custom struct as an argument when launching CubeCL kernels.I was not entirely sure where to place all of my code, here is a summary of my approach:
cubecl-core, I added a new traitAsLaunchArgumentdefining the signature of theas_launch_arg()functioncubecl-macros, I added a new macro calledCubeLaunchArgthat implementsAsLaunchArgumentfor a structTypeField::is_scalar(&Type)for this); if so, wrap the field inScalarArg, if not, hope that the field implementsAsLaunchArg.This is very rudimentary, and is just enough to make my use case work. I would be happy to extend it further, but I think it makes sense to check out whether my approach makes sense to you (as in the CubeCL maintainers & devs).
Thanks!
Marcel