heterocl.tvm.tensor_intrin module

Tensor intrinsics

class TensorIntrin(handle)

Bases: heterocl.tvm._ffi.node.NodeBase

Tensor intrinsic functions for certain computation.

See also

decl_tensor_intrin

Construct a TensorIntrin

decl_tensor_intrin(op, fcompute, name='tensor_intrin', binds=None)

Declare a tensor intrinsic function.

Parameters

• op (Operation) – The symbolic description of the intrinsic operation

• fcompute (lambda function of inputs, outputs-> stmt) –

  Specifies the IR statement to do the computation. See the following note for function signature of fcompute

  Note

  Parameters

  • ins (list of Buffer) - Placeholder for each inputs

  • outs (list of Buffer) - Placeholder for each outputs

  Returns

  • stmt (Stmt, or tuple of three stmts)

  • If a single stmt is returned, it represents the body

  • If tuple of three stmts are returned they corresponds to body, reduce_init, reduce_update

• name (str, optional) – The name of the intrinsic.

• binds (dict of Tensor to 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.

Returns

intrin – A TensorIntrin that can be used in tensorize schedule.

Return type

TensorIntrin