AI-generated implementations for DaCe library nodes - #2557
Conversation
7e32b3a to
9943722
Compare
| node to begin with -- a microkernel DaCe does not expose, a vendor intrinsic sequence, a call | ||
| into an external library -- so the description *is* the specification:: | ||
|
|
||
| node = AINode('fma', 'Compute out = a * b + c elementwise over the whole tile.', |
There was a problem hiding this comment.
I think the pure/naive lowering of the Library can be provided to describe the semantics on a more constrained way?
There was a problem hiding this comment.
the pure/native lowering can be arbitrarily complex (with e.g., multiple maps in reduce) and so it might require a lot of prompt real estate and to actually use the dace codegen (or print out e.g. a schedule tree?)
Not sure. Potential future work?
There was a problem hiding this comment.
Makes sense, also to consider a library node implementation valid we should use symbolic shapes to create fuzzed correctness tests.
I think this should be a future work, I'm doing this for the ICLR submission, we can expand this after ICLR submission. Having measures against reward-hacking and vagueness is critical to avoid token wastage.
| desc='Natural-language description of what this node must compute. This is ' | ||
| 'the specification handed to the model, so state the intended semantics, ' | ||
| 'the expected numerical behavior, and any implementation technique that is ' | ||
| 'required (e.g. "use AVX2 intrinsics").') |
There was a problem hiding this comment.
I think DaCe could (or should?) provide a default pre- or post-amble for the description that DaCe requires input language in C++, CUDA, HIP or Python for library node implementations.
There was a problem hiding this comment.
Yes, already part of the prompt (see prompt.py)
…bject, and default schedule/storage inference
Updated pytest commands to exclude tests marked with 'ai'.
2 participants
This PR provides two elements:
'ai'implementation, available on every library node. Expanding with it asks a language model to write the replacement tasklet, given the node's context in the SDFG and on the local machine.dace.libraries.ai.nodes.ai_node.AINode, a library node that carries only a natural-language description, for microkernels that have no library node of their own.ai.refineallows AI sessions to refine the implementation (e.g., if it crashes or not fast enough)The model is reached through a provider selected by the
ai.providerconfiguration entry. It supports the Anthropic Messages API, the OpenAI Responses API, and arbitrary AI models through amanualmode that lets you paste prompts and responses.