Using compile targets to customize compilation

Note: Support for compile targets is currently experimental and may change in the future.

Compile targets can be used to describe multiple stages of compilation, code generation, as well as to infer compilation pipelines from target capabilities. They can be defined using the cudaq._experimental.CompileTarget Python class.

Constructing single-stage compile targets

Compile targets composed of a single compilation stage can easily be constructed using one of the following constructors:

• CompileTarget(ir_output_profile=profile) defines a target made of a single compilation stage that lowers the input kernel to the specified IR profile. The supported profiles are QuakeReferenceSemanticsProfile and QuakeUnmappedProfile.

• CompileTarget(custom_pipeline=lowering) defines a target made of a single compilation stage, given by the custom lowering pipeline. lowering should be either a CustomPipelineLowering or a MLIR pipeline string.

More complex multi-stage compile targets can be constructed using the CompileTarget(stages=[]) constructor, although support for this is currently limited. A QuakeMappedProfile stage routes the kernel onto its device when one is set. stages, ir_output_profile and custom_pipeline are mutually exclusive.

Setting compile targets with set_target

To use a compile target for all kernel launches that follow, use the set_target function with the target name and the compile target:

• Use cudaq.set_target("name", compile_target=target) to run kernels against a pre-defined named CUDA-Q target, but overriding the compile target.

• Use cudaq.set_target(CustomTarget(compile_target=target, runtime_endpoint=endpoint)) to run kernels using a custom compile target and a custom runtime endpoint.

Example

The example below builds an inferred target and an explicit-pipeline target, then compiles a kernel against each through the real set_target path. A custom runtime endpoint is used to capture the lowered MLIR rather than executing it.

from contextlib import contextmanager

import cudaq
from cudaq._experimental import (
    CompileTarget,
    CustomPipelineLowering,
    CustomTarget,
    QuakeReferenceSemanticsProfile,
)

basis = ["h", "s", "t", "rx", "ry", "rz", "x", "y", "z", "x(1)"]

# Define a custom target, with the pipeline inferred by the compiler.
inferred = CompileTarget(
    ir_output_profile=QuakeReferenceSemanticsProfile(
        basis_gates=basis, measurement_type="single_qubit"),
    codegen="qir-base",
)
# The target has one compilation stage, with implicit lowering.
assert len(inferred.stages) == 1
assert inferred.stages[0].lowering is None

# Alternatively, a target can also be defined by specifying the compilation
# pipeline explicitly.
explicit = CompileTarget(
    custom_pipeline=CustomPipelineLowering(
        high_level="expand-measurements",
        mid_level="basis-conversion{basis=" + ",".join(basis) + "}",
    ),
    codegen="qir-base",
)
# The target has one compilation stage, with an explicit custom lowering.
assert len(explicit.stages) == 1
assert explicit.stages[0].lowering is not None


@cudaq.kernel
def bell():
    q = cudaq.qvector(2)
    h(q[0])
    x.ctrl(q[0], q[1])
    mz(q)


@contextmanager
def set_target_context(*args, **kwargs):
    """Set the target context for the duration of the context manager."""
    cudaq.set_target(*args, **kwargs)
    try:
        yield
    finally:
        cudaq.reset_target()


# Run the kernel on an existing target, using the specified compile pipeline.
with set_target_context("qpp-cpu", compile_target=explicit):
    res = cudaq.sample(bell, shots_count=1)
    assert "00" in res or "11" in res


# Or alternatively, define your custom runtime endpoint:
class InspectEndpoint:
    """Captures the lowered MLIR instead of executing the kernel."""

    def sample(self, module, args, **kwargs):
        self.mlir = str(module.mlir_module)
        return cudaq.SampleResult()


# and customize both the compile target and runtime endpoint:
endpoint = InspectEndpoint()
target = CustomTarget(compile_target=inferred, runtime_endpoint=endpoint)
with set_target_context(target):
    res = cudaq.sample(bell, shots_count=1)
    assert "quake.mz" in endpoint.mlir
    print(endpoint.mlir)