hidet.drivers

hidet.drivers

Functions:

build_ir_module(ir_module, output_dir, target)

Build an IR module to a shared library or object file.

build_ir_module_batch(ir_modules, ...[, force])

Build a batch of IR modules.

build_task(task[, target, load])

Build a task into a compiled function.
hidet.drivers.build_ir_module(ir_module, output_dir, target, output_kind='.so', force=False)[source]

Build an IR module to a shared library or object file.

This driver function performs the following steps to build an IR module: 1. Lower and optimize the IR module with a sequence of pre-defined passes. 2. Generate source code from the lowered IR module. 3. Call the underlying compiler (e.g., gcc or nvcc) to compile the generated source code into a shared library

(when output_kind == ‘.so’) or an object file (when output_kind == ‘.o’).

To ensure safe parallel execution in a multiprocessing environment, a file-based lock (.lock file in the output_dir) is used. This guarantees that only one process can build the IR module for a given output_dir at any given time.

Parameters:

  • ir_module (Union[IRModule, Sequence[IRModule]]) – The IR module to be built. This can be a single IRModule or a sequence of IRModules.

  • output_dir (str) – The directory to save the generated source code and the compiled library.

  • target (str) – The target to build the IR module. Supported targets are cpu and cuda. Attributes (e.g., ‘cuda –arch=sm_70’) can also be specified.

  • output_kind (str) – The output kind. Supported kinds are ‘.so’ and ‘.o’. - ‘.so’: Compile the IR module to a shared library. - ‘.o’: Compile the IR module to an object file.

  • force (bool) – Whether to force re-build the IR module. By default, the IR module will not be re-built if the library already exists in the specified output directory.

Notes

  • File Locking: A .lock file is created in the output_dir to synchronize access. If another process tries to build the same IR module concurrently, it will wait until the lock is released.

  • Parallel Safety: The file lock ensures that only one process builds the IR module for a specific output_dir.

hidet.drivers.build_ir_module_batch(ir_modules, output_dirs, output_kind, target, force=False)[source]

Build a batch of IR modules.

Parameters:

  • ir_modules (Sequence[IRModule]) – A sequence of IR modules to build.

  • output_dirs (Sequence[str]) – Directories for compilation artifacts.

  • output_kind (str) – The output kind of the compiled library. Can be ‘.so’ or ‘.o’.

  • target (str) – The target of the compilation. Can be ‘cuda’ or ‘cpu’.

  • force (bool) – Whether to force re-build the IR module. By default, the IR module will not be re-built if the library already exists in the specified output directory.

hidet.drivers.build_task(task, target='cuda', load=True)[source]

Build a task into a compiled function.

Parameters:

  • task (Task) – The task to be built.

  • target (str) – The target platform. Candidates are ‘cuda’ and ‘cpu’.

  • load (bool) – Whether to load the compiled function. If False, the compiled function will not be loaded, and None is returned. Otherwise, the compiled function is loaded and returned.

Returns:

ret – When load is True, the compiled function is returned. Otherwise, None is returned.

Return type:

CompiledTask