In this example, we will write a program in hidet script to print out "hello, world".
Before we start, we need to install two python packages with the following cell to enable stdout capture.
%%capture
%load_ext wurlitzer
import hidet
from hidet.lang import printf, attr
# declare a script module, all script function inside the with-context will be added to it
with hidet.script_module() as script_module:
# all functions decorated by "hidet.script" will be marked as Hidet Script function
@hidet.script
def main():
# mark this function is a host side function (run on cpu)
attr.func_kind = 'host_kernel'
# print out
printf('hello, world!\n')
# get the ir module (hidet.ir.IRModule), which is the top-level IR node
ir_module = script_module.ir_module()
print('IRModule:')
print(str(ir_module).strip()) # hidet.ir.IRModule
print()
IRModule: def main fn() # func_kind: host_kernel printf("hello, world!\n");
# compile the ir module to a dynamic library and load it into a python callable
compiled_func = hidet.driver.build_ir_module(ir_module, func_name='main')
print(' Source path: "{}"'.format(compiled_func.src_path))
print(' Compiled library: "{}"'.format(compiled_func.lib_path))
Source path: "./outs/ir_module/1ba570cf9a404d0f/source.cu" Compiled library: "./outs/ir_module/1ba570cf9a404d0f/lib.so"
Hidet will compile an IRModule to a dynamic library (e.g., lib.so
), and dynamically load it into the python environment.
print('Source code:\n{}'.format(compiled_func.source(color=True)))
Source code: #include <stdint.h> #include <cuda_fp16.h> #include <cuda_bf16.h> #include <hidet/runtime/cuda_context.h> #include <hidet/runtime/cpu_context.h> typedef float tfloat32_t; #define __float_to_tf32(x) (x) extern "C" { __host__ void hidet_main() { printf("hello, world!\n"); } }
print('Run the compiled program:')
compiled_func()
Run the compiled program: hello, world!