Skip to content

Use hir

Ryan Ziegler requested to merge ryanjz2/hercules-onnx:use-hir into main

Create a use_hir macro, which enables embedding Hercules IR into a macro and running functions defined within. This roughly works as follows:

  • Generate an object file with a unique name for each invocation of the macro
  • Link against this object file by providing its name to codegen, which has been modified to allow for generating linking properties for the Rust compiler

Note that we must do the second step this way because build scripts run before macro evaluation. I've tested this with a simple matmul, shown below:

hercules_rt::use_hir!(
"
fn matmul<3>(a: array(f32, #0, #1), b: array(f32, #1, #2)) -> array(f32, #0, #2)
  c = constant(array(f32, #0, #2), [])
  i_j_ctrl = fork(start, #0, #2)
  i_idx = thread_id(i_j_ctrl, 0)
  j_idx = thread_id(i_j_ctrl, 1)
  k_ctrl = fork(i_j_ctrl, #1)
  k_idx = thread_id(k_ctrl, 0)
  k_join_ctrl = join(k_ctrl)
  i_j_join_ctrl = join(k_join_ctrl)
  r = return(i_j_join_ctrl, update_i_j_c)
  zero = constant(f32, 0)
  a_val = read(a, position(i_idx, k_idx))
  b_val = read(b, position(k_idx, j_idx))
  mul = mul(a_val, b_val)
  add = add(mul, dot)
  dot = reduce(k_join_ctrl, zero, add)
  update_c = write(update_i_j_c, dot, position(i_idx, j_idx))
  update_i_j_c = reduce(i_j_join_ctrl, c, update_c)
"
);
Edited by Ryan Ziegler

Merge request reports

Loading