tvm-clj.compiler

Once a user has a schedule, they then need to compile the schedule into actual executable code. This produces a module which contains all of the executable code, cuda/opencl modules, etc required to run the schedule.

compile

(compile fn-map {:keys [target-host], :or {target-host "llvm"}})(compile fn-map)

Build a map of function entries. fn-entries is a map of name to fn-data and fn-data is a map containing:

  • :schedule - schedule to use.
  • :arguments - argument declarations to the function
  • :bind-map - optional map of argument to bind declaration to declare the memory layout of the argument.

int->lowered-function-type-map

ir->fn

(ir->fn ir-data fn-name)

Given map of {:schedule :arguments} containing optionally :target create a clojure fn that calls into TVM.

lower

(lower sch args {:keys [name binds simple-mode? optimization-level], :or {name "main", optimization-level 2}})

Lowering step before build into target.

Users do not often need to call this themselves but it is useful in order to see code produced by different scheduling primitives. The return value prints the schedule nicely to the REPL.

Parameters ———- sch : tvm.te.schedule.Schedule The schedule to be built

args : list of Buffer or Tensor or Var The argument lists to the function.

name : str, optional The name of result function.

binds : dict of :any:Tensor to :any:Buffer, optional Dictionary that maps the Tensor to Buffer which specified the data layout requirement of the function. By default, a new compact buffer is created for each tensor in the argument.

simple_mode : bool, optional Whether only output simple and compact statement, this will skip LoopPartition, api wrapper generation and Unrolling.

Returns ——- m : IRModule or Stmt The result IRModule, if simple_mode=False Then the Stmt before make api is returned.

lowered-function-type->int-map

view-ir

(view-ir ir-data)

Lower the schedule without optimizations returning an IR layer that has a nice pprint for inspecting schedule state.