AI-generated implementations for DaCe library nodes

.github/workflows/general-ci.yml

@@ -65,7 +65,7 @@ jobs:
         else
             export DACE_optimizer_automatic_simplification=${{ matrix.simplify }}
         fi
-        uv run  --group testing pytest -n auto --cov-report=xml --cov=dace --tb=short --timeout_method thread --timeout=300 -m "not gpu and not autodiff and not torch and not onnx and not tensorflow and not mkl and not sve and not papi and not mlir and not lapack and not mpi and not scalapack and not datainstrument and not long and not sequential"
+        uv run  --group testing pytest -n auto --cov-report=xml --cov=dace --tb=short --timeout_method thread --timeout=300 -m "not gpu and not autodiff and not torch and not onnx and not tensorflow and not mkl and not sve and not papi and not mlir and not lapack and not mpi and not scalapack and not datainstrument and not long and not sequential and not ai"
         ./codecov
 
     - name: Test OpenBLAS LAPACK

.github/workflows/gpu-ci.yml

@@ -56,13 +56,13 @@ jobs:
         source ~/.venv/bin/activate # activate venv
         export PATH=$PATH:/usr/local/cuda/bin  # some test is calling cuobjdump, so it needs to be in path
         echo "CUDACXX: $CUDACXX"
-        pytest --cov-report=xml --cov=dace --tb=short --timeout_method thread --timeout=300 -m "gpu and not mpi"
+        pytest --cov-report=xml --cov=dace --tb=short --timeout_method thread --timeout=300 -m "gpu and not mpi and not ai"
 
     - name: Run pytest GPU with distributed compilation
       run: |
         source ~/.venv/bin/activate # activate venv
         export PATH=$PATH:/usr/local/cuda/bin
-        mpirun -n 2 --oversubscribe pytest --tb=short --timeout_method thread --timeout=300 --with-mpi -m "gpu and mpi"
+        mpirun -n 2 --oversubscribe pytest --tb=short --timeout_method thread --timeout=300 --with-mpi -m "gpu and mpi and not ai"
 
     - name: Run extra GPU tests
       run: |

ci/cscs_gpu.yml

@@ -36,7 +36,7 @@ test_cscs_gh200:
     - .container-runner-daint-gh200
   image: ${CSCS_REGISTRY_PATH}/public/dace-gpu:${DOCKER_TAG}
   variables:
-    TEST_VARIANTS: 'gpu and not mpi'  # no MPI launcher here; the GPU workflow covers those
+    TEST_VARIANTS: 'gpu and not mpi and not ai'  # no MPI launcher here; the GPU workflow covers those
     SLURM_GPUS_PER_NODE: 4
     SLURM_JOB_NUM_NODES: 1
     SLURM_TIMELIMIT: 30

dace/codegen/targets/cpu.py

@@ -1584,6 +1584,12 @@ def _generate_Tasklet(self,
         self._locals.clear_scope(self._ldepth + 1)
         self._dispatcher.defined_vars.exit_scope(node)
 
+    #: ``generate_node`` dispatches on the exact class name, so a tasklet subclass needs an entry of
+    #: its own. An AI-generated tasklet is a tasklet in every respect the code generator cares
+    #: about; it differs only in recording the library node and conversation it came from, which is
+    #: read back by :mod:`dace.libraries.ai.iterate`.
+    _generate_AITasklet = _generate_Tasklet
+
     def unparse_tasklet(self, sdfg, cfg, state_id, dfg, node, function_stream, inner_stream, locals, ldepth,
                         toplevel_schedule):
         # Call the generic CPP unparse_tasklet method

dace/config_schema.yml

@@ -584,6 +584,141 @@ required:
                         default_Darwin: ''
                         default_Windows: ''
 
+    ai:
+        type: dict
+        title: AI-generated code
+        description: >
+            Preferences for the "ai" library node implementation, which expands a library node
+            into a tasklet written by a language model.
+        required:
+            provider:
+                type: str
+                default: 'anthropic'
+                title: Model provider
+                description: >
+                    Which backend to generate code with. "anthropic" uses the Anthropic Messages
+                    API (pip install 'dace[ai-anthropic]'); "responses" uses the Responses API of
+                    the openai package (pip install 'dace[ai-openai]'); "manual" prints the prompt and
+                    reads the reply back from the terminal, needing neither a key nor an SDK. The
+                    SDK is only imported when an expansion actually runs.
+
+            manual_dir:
+                type: str
+                default: ''
+                title: Manual provider directory
+                description: >
+                    Directory the "manual" provider writes prompts to and reads replies from.
+                    Defaults to a "dace_ai_prompts" directory under the system temporary directory
+                    when empty.
+
+            api_key_envvar:
+                type: str
+                default: 'ANTHROPIC_API_KEY'
+                title: API key environment variable
+                description: >
+                    Name of the environment variable holding the provider's API key. If
+                    the variable is unset, the provider's own credential chain is used instead.
+
+            model:
+                type: str
+                default: 'claude-opus-5'
+                title: Model
+                description: Identifier of the model used to generate tasklets.
+
+            effort:
+                type: str
+                default: 'high'
+                title: Reasoning effort
+                description: >
+                    How much effort the model should spend per expansion, for providers that
+                    support it (low, medium, high, xhigh, max).
+
+            max_tokens:
+                type: int
+                default: 64000
+                title: Maximum response tokens
+                description: Upper bound on the length of a single generated response.
+
+            timeout:
+                type: int
+                default: 600
+                title: Request timeout
+                description: Timeout, in seconds, for one request to the model provider.
+
+            verify:
+                type: bool
+                default: true
+                title: Verify generated code
+                description: >
+                    Compile the generated tasklet in isolation before it is inserted into the
+                    SDFG, and send any compiler diagnostics back to the model. This catches
+                    syntax and type errors; it does not check that the result is correct.
+
+            max_repair_attempts:
+                type: int
+                default: 2
+                title: Maximum repair attempts
+                description: >
+                    How many times the model may be asked to fix generated code that does not
+                    compile. Only used when verification is enabled.
+
+            environment_dir:
+                type: str
+                default: ''
+                title: Generated environment directory
+                description: >
+                    Directory where library environments requested by the model are written as
+                    Python modules, so that they can be inspected, edited, and resolved again in
+                    later sessions. Defaults to ~/.dace/ai_environments when empty.
+
+            cache:
+                type: bool
+                default: true
+                title: Reuse model answers
+                description: >
+                    Store every answer on disk, keyed by the provider, the model, the effort
+                    setting and the whole conversation, and reuse it when the same question is
+                    asked again. Re-expanding an unchanged library node then costs nothing. Any
+                    change to the SDFG, the node, or the model changes the key and misses the
+                    cache, so a stale answer cannot be picked up.
+
+            cache_dir:
+                type: str
+                default: ''
+                title: Answer cache directory
+                description: >
+                    Where reusable answers are stored. Defaults to ~/.dace/ai_cache when empty.
+                    Deleting the directory, or any file in it, only forces the question to be
+                    asked again.
+
+            sessions:
+                type: bool
+                default: true
+                title: Record generation sessions
+                description: >
+                    Write the prompts, the model's answers, and the source and diagnostics of every
+                    probe compilation to disk, one directory per generated slot. An expansion is a
+                    paid and non-reproducible call whose result is baked into the SDFG, so the
+                    material that produced it is kept by default. It is also what a later round
+                    resumes: refining a tasklet replays this conversation rather than starting
+                    over. Set DACE_debugprint=verbose to see the same material as it happens.
+
+            session_dir:
+                type: str
+                default: ''
+                title: Session directory
+                description: >
+                    Where generation sessions are written. Defaults to ~/.dace/ai_sessions when
+                    empty.
+
+            extra_instructions:
+                type: str
+                default: ''
+                title: Extra prompt instructions
+                description: >
+                    Additional text appended to every code generation prompt, e.g. project-wide
+                    coding conventions or libraries that should be preferred.
+
     instrumentation:
         type: dict
         title: Instrumentation

dace/frontend/python/interface.py

@@ -437,6 +437,61 @@ def in_program() -> bool:
     return False
 
 
+def ai(description: str, *inputs, **kwargs):
+    """
+    Computes something described in natural language, rather than written out.
+
+    The call becomes an :class:`~dace.libraries.ai.nodes.ai_node.AINode` in the SDFG: a library node
+    that carries the description instead of an implementation. It is expanded when the program is
+    compiled, by asking a language model to write the tasklet that goes in its place, given the
+    description and the node's context -- the connectors and the data behind them, the maps and
+    nested SDFGs around it, and the compiler and machine it will be built for. The generated
+    tasklet is then part of the SDFG like any other node: it is saved with it, and compiling the
+    program again does not query the model again.
+
+    Use it where the implementation is the point and the code is not worth writing by hand: a
+    vectorized microkernel, a vendor intrinsic sequence, a call into an external library::
+
+        @dace.program
+        def matmul(A: dace.float32[N, N], B: dace.float32[N, N]):
+            return dace.ai('Write _out[i][j] = sum over k of _a[i][k] * _b[k][j], as a '
+                           'register-blocked microkernel using AVX2 FMA intrinsics.',
+                           a=A, b=B)
+
+    Since the description is the specification, it has to be able to name the variables the
+    generated code will see. Those are the node's connectors, and they are named from the call:
+    a keyword input ``a=A`` becomes ``_a``, a positional input becomes ``_in0``, ``_in1``, ... in
+    order, and the output is ``_out`` (or ``_out0``, ``_out1``, ... when ``out`` names several).
+    Symbols of the program, such as ``N`` above, are in scope in the generated code under their own
+    names and do not have to be passed in.
+
+    Which model is asked, and whether its answers are cached, is set by the ``ai`` entries of the
+    DaCe configuration -- see :mod:`dace.libraries.ai`. To read what the model wrote, or to ask it
+    for something better, see :func:`dace.libraries.ai.refine`.
+
+    :param description: What the node must compute, in natural language. State the intended
+                        semantics, the expected numerical behavior, and any implementation
+                        technique that is required (e.g. "use AVX2 intrinsics").
+    :param inputs: Data containers to read, connected as ``_in0``, ``_in1``, ...
+    :param out: Data container to write to, or a sequence of them. If not given, one is allocated
+                and returned. It has to name a whole container: to write into part of one, assign
+                the result instead, as in ``C[0:8, 0:8] = dace.ai(...)``.
+    :param shape: Shape of the allocated output. Defaults to the shape of the first input.
+    :param dtype: Data type of the allocated output. Defaults to the type of the first input.
+    :param storage: Storage type of the allocated output. Defaults to that of the first input.
+    :param name: Name of the node, shown in the SDFG and given to the model with the description.
+    :param schedule: Schedule of the node, which determines its default device mapping.
+    :param kwargs: Any other keyword argument is a data container to read, connected under a
+                   connector named after the keyword (``a=A`` is read as ``_a``).
+    :return: The container written to, or a tuple of them if ``out`` named several.
+    :note: This function is parsed by the DaCe Python frontend and cannot be called outside of a
+           ``@dace.program``. Use :class:`~dace.libraries.ai.nodes.ai_node.AINode` directly to add
+           such a node to an SDFG built with the SDFG API.
+    """
+    raise NotImplementedError('dace.ai can only be called inside a DaCe program. To describe a node in an '
+                              'SDFG built with the SDFG API, add a dace.libraries.ai.AINode to a state.')
+
+
 class named:
     """
     Creates a `NamedRegion` with the given label.

dace/frontend/python/replacements/__init__.py

@@ -3,6 +3,7 @@
 from .utils import *
 
 # ``from x import *`` is used to not break compatibility with previous versions of dace
+from .ai import *
 from .array_creation import *
 from .array_creation_dace import *
 from .array_creation_cupy import *