Build System

def entrypoint(f=None, *, description: str | None = None):
    """
    Function decorator to turn it into a build entrypoint.
    
    Parameters:
    - f: Function to decorate
    - description: Optional description for the entrypoint
    
    Returns:
    Entrypoint: Build entrypoint object
    """

def cl_arg(name: str, **kwargs):
    """
    Decorator for defining command-line arguments.
    
    Parameters:
    - name: Argument name
    - **kwargs: Additional argument parser configuration
    """

class BuildAction(BuildDependency):
    """
    Base class for build system actions.
    
    Actions represent individual build steps that can be composed
    into complex workflows with dependency management. This class is
    designed to be subclassed by concrete actions.
    """
    
    def _invoke(self) -> None:
        """
        Execute the action in-process (override for in-process actions).
        """
    
    def _remotable_thunk(self):
        """
        Return a remotable thunk for out-of-process execution.
        
        Returns:
        Tuple of (callable, args) for out-of-process execution
        """

class CompileAction(BuildAction):
    """
    Action for compiling IREE modules from source.
    """
    
    def __init__(self, inv: CompilerInvocation, **kwargs): ...

class ImportOnnxAction(BuildAction):
    """
    Action for importing ONNX models to MLIR.
    """
    
    def __init__(self, input_file: BuildFile, output_file: BuildFile, upgrade: bool, **kwargs): ...

class FetchHttpAction(BuildAction):
    """
    Action for fetching resources over HTTP.
    """
    
    def __init__(self, url: str, output_file: BuildFile, **kwargs): ...

class ExecuteOutOfProcessThunkAction(BuildAction):
    """
    Executes a callback thunk with arguments out-of-process.
    
    Both the thunk and args must be pickleable.
    """

class Entrypoint:
    """
    Build system entrypoint wrapping a callable with argument definitions.
    """
    
    def __init__(self, name: str, wrapped: Callable, description: str | None = None): ...
    
    def __call__(self, *args, **kwargs): ...

class BuildFile:
    """
    Represents a file in the build system with dependency tracking.
    """
    
    def __init__(self, namespace: FileNamespace, path: str): ...

class FileNamespace(Enum):
    """
    File namespace enumeration for organizing build artifacts.
    """
    GEN = "gen"  # Transient generated files
    PARAMS = "params"  # Distributable parameter files
    BIN = "bin"  # Distributable, platform-neutral binaries
    PLATFORM_BIN = "platform_bin"  # Platform-specific binaries

class BuildContext:
    """
    Build execution context managing file namespaces and dependencies.
    """
    
    @staticmethod
    def current() -> BuildContext: ...
    
    def gen_file(self, path: str) -> BuildFile: ...
    
    def bin_file(self, path: str) -> BuildFile: ...

class ActionConcurrency(Enum):
    """
    Action concurrency execution models.
    """
    THREAD = "thread"  # Thread-based concurrency
    PROCESS = "process"  # Process-based concurrency
    NONE = "none"  # Sequential execution

def compile(
    input_file: BuildFileLike,
    *,
    target_backends: Sequence[str],
    output_file: BuildFile | None = None,
    **kwargs
) -> BuildFile:
    """
    Compile an MLIR source file to VM bytecode using the build system.
    
    Parameters:
    - input_file: Source file to compile
    - target_backends: List of target backends
    - output_file: Optional output file (auto-generated if not provided)
    - **kwargs: Additional compiler options
    
    Returns:
    BuildFile: Output file containing compiled bytecode
    """

def expand_cl_arg_defaults(func: Callable) -> Callable:
    """
    Expand command-line argument defaults for a function.
    
    Parameters:
    - func: Function to process
    
    Returns:
    Callable: Function with expanded defaults
    """

def extract_cl_arg_defs(func: Callable) -> List:
    """
    Extract command-line argument definitions from a function.
    
    Parameters:
    - func: Function to analyze
    
    Returns:
    List: Command-line argument definitions
    """

def current_args_namespace():
    """
    Get the current command-line arguments namespace.
    
    Returns:
    argparse.Namespace: Current arguments
    """

import iree.build
from iree.build import entrypoint, cl_arg

# Define a simple build entrypoint
@entrypoint(description="Compile MLIR model")
@cl_arg("--input", required=True, help="Input MLIR file")
@cl_arg("--output", help="Output VM file")
@cl_arg("--target", default="llvm-cpu", help="Target backend")
def compile_model(input_file: str, output_file: str = None, target: str = "llvm-cpu"):
    """Simple model compilation entrypoint."""
    context = iree.build.BuildContext.current()
    
    # Create input and output files
    input_build_file = context.gen_file(input_file)
    if output_file is None:
        output_file = input_file.replace(".mlir", ".vmfb")
    output_build_file = context.bin_file(output_file)
    
    # Compile using build system
    result = iree.build.compile(
        input_build_file,
        target_backends=[target],
        output_file=output_build_file
    )
    
    return result

from iree.build import entrypoint, cl_arg, BuildContext
from iree.build import ImportOnnxAction, CompileAction

@entrypoint(description="Import and compile ONNX model")
@cl_arg("--onnx-file", required=True, help="Input ONNX file")
@cl_arg("--targets", nargs="+", default=["llvm-cpu"], help="Target backends")
@cl_arg("--upgrade", action="store_true", help="Upgrade ONNX model")
def compile_onnx_model(onnx_file: str, targets: list, upgrade: bool = False):
    """Import ONNX model and compile for multiple targets."""
    context = BuildContext.current()
    
    # Set up files
    input_file = context.gen_file(onnx_file)
    mlir_file = context.gen_file(onnx_file.replace(".onnx", ".mlir"))
    
    # Import ONNX to MLIR
    import_action = ImportOnnxAction(
        input_file=input_file,
        output_file=mlir_file,
        upgrade=upgrade
    )
    
    # Compile for each target
    results = []
    for target in targets:
        output_file = context.bin_file(f"model_{target}.vmfb")
        
        compile_result = iree.build.compile(
            mlir_file,
            target_backends=[target],
            output_file=output_file
        )
        results.append(compile_result)
    
    return results

from iree.build import entrypoint, cl_arg, BuildContext
from iree.build import FetchHttpAction, compile

@entrypoint(description="Fetch model from URL and compile")
@cl_arg("--url", required=True, help="Model URL")
@cl_arg("--target", default="llvm-cpu", help="Target backend")
def fetch_and_compile(url: str, target: str = "llvm-cpu"):
    """Fetch a model from URL and compile it."""
    context = BuildContext.current()
    
    # Determine file extension from URL
    if url.endswith(".onnx"):
        local_file = context.gen_file("downloaded_model.onnx")
    else:
        local_file = context.gen_file("downloaded_model.mlir")
    
    # Fetch the file
    fetch_action = FetchHttpAction(
        url=url,
        output_file=local_file
    )
    
    # Compile the downloaded file
    output_file = context.bin_file(f"model_{target}.vmfb")
    result = compile(
        local_file,
        target_backends=[target],
        output_file=output_file
    )
    
    return result

from iree.build import entrypoint, cl_arg, BuildAction, BuildContext

class CustomPreprocessAction(BuildAction):
    """Custom preprocessing action."""
    
    def __init__(self, input_file, output_file, **kwargs):
        super().__init__(**kwargs)
        self.input_file = input_file
        self.output_file = output_file
        self.output_file.deps.add(self)
        self.deps.add(self.input_file)
    
    def _invoke(self):
        # Custom preprocessing logic
        with open(self.input_file.get_fs_path(), 'r') as f:
            content = f.read()
        
        # Apply preprocessing
        processed_content = content.replace("old_pattern", "new_pattern")
        
        with open(self.output_file.get_fs_path(), 'w') as f:
            f.write(processed_content)

@entrypoint(description="Custom preprocessing and compilation")
@cl_arg("--input", required=True, help="Input file")
@cl_arg("--target", default="llvm-cpu", help="Target backend")
def custom_build(input_file: str, target: str = "llvm-cpu"):
    """Custom build with preprocessing step."""
    context = BuildContext.current()
    
    # Input and intermediate files
    raw_input = context.gen_file(input_file)
    processed_input = context.gen_file("processed_" + input_file)
    
    # Custom preprocessing
    preprocess_action = CustomPreprocessAction(
        input_file=raw_input,
        output_file=processed_input
    )
    
    # Compile preprocessed file
    output_file = context.bin_file(f"custom_model_{target}.vmfb")
    result = compile(
        processed_input,
        target_backends=[target],
        output_file=output_file
    )
    
    return result