Open int3 opened 1 month ago
Another approach would be to write a simple ast transform that inserts these _builder arguments
I'm actually not sure if this is correct / safe; are we guaranteed that the function will get JIT-compiled? Or do we have to actually mark it with @jit
? Looking at the MLIR, it looks like things get JIT-compiled regardless of whether we add @jit
, but I'm not sure if it's just working by accident.
I've been thinking about how to support some libdevice operations that don't map cleanly to existing vector math libraries on the CPU. It seems to me that things like
isnan
could be implemented reasonably cleanly in Triton itself. E.g. for the fp32 case:However, libdevice.isnan as defined in the interface file triton/language/extra/libdevice.py is not marked as
@jit
but rather as @core.extern. So we can't use the above syntax; instead we need to explicitly specify the _builder argument:which is pretty ugly.
So I'm wondering if
@core.extern
vs@jit
should be part of the libdevice interface. On one hand, it seems like an implementation detail; on the other hand, the "calling convention" of a function is traditionally part of its interface. I think making it into part of the implementation is quite doable; we would simply have to replace the invoke-timedispatch
function with something that does the mapping atlibdevice
-module-creation time.Another approach would be to write a simple
ast
transform that inserts these_builder
arguments, so if we wanted to declare anextern
function with an implicit builder, it would be something likeThis might be simpler / less invasive. Would love to hear your thoughts.