Struct OpDef

pub struct OpDef(/* private fields */);

C++ OpDef: register a new operator from Rust.

Each step updates the registry at once, as in C++. After a failed step the rest are skipped. The chain completes when it is dropped, at the end of its statement, as a C++ OpDef does: the operator’s APIs are requested in registered destinations, and the error of a failed step panics, which a tvm_ffi::static_init_block! reports as the error of its block. OpDef::op completes it instead, and returns the operator or the error. The registry does not synchronize registration with readers, so register an operator before building IR with it, as a tvm_ffi::static_init_block! does. An operator and its hooks live for the rest of the process.

use tvm_tirx::ir::op_attr::CALL_EFFECT_KIND;
use tvm_tirx::ir::sig::{self, IntExpr};
use tvm_tirx::ir::{CallEffectKind, CallObj, OpDef, Type};
use tvm_tirx::tvm_ffi::{self, Result};

fn first_arg_type(call: &CallObj) -> Result<Type> {
    Ok(call.args.get(0)?.ty.clone())
}

tvm_ffi::static_init_block! {
    OpDef::new("my_ext.clamp", "Clamp an index to a bound.")
        .signature((sig::arg::<IntExpr>("index", ""), sig::arg::<IntExpr>("bound", "")))
        .set_infer_type(first_arg_type)
        .set_attr(CALL_EFFECT_KIND, CallEffectKind::Pure);
}

Implementations

impl OpDef

pub fn new(name: &str, doc: &str) -> Self

Register a new operator (ir.RegisterOp), with doc as its documentation, less the indentation its lines share and the blank lines around it, so that a string of several lines can be indented with the code. A name already registered is an error.

pub fn op(self) -> Result<Op>

Complete the chain now: the operator, or the error of the step that failed, which dropping the chain would panic with.

pub fn signature<S: Specs>(self, specs: S) -> Self

C++ OpDef::signature: replace the argument descriptions and the attrs type, and the validator generated from them, which checks a Call’s argument counts, argument types and attrs. specs is a descriptor of sig, or a tuple of them in order, as C++ takes them. A custom validator keeps precedence. Without sig::call_attrs the operator has no attrs type, and Call.attrs is unconstrained.

pub fn add_arg(self, name: &str, doc: &str) -> Self

Append a value argument’s description, without changing what the validator checks.

pub fn add_ty_arg(self, name: &str, doc: &str) -> Self

Append a type argument’s description, without changing what the validator checks.

pub fn attrs_type_key(self, key: &str) -> Self

The type of Call.attrs, by its registered type key.

pub fn set_attr<T: Into<Any>>( self, key: AttrKey<T>, value: impl Into<T>, ) -> Self

Register the attribute key (C++ set_attr(key, value)); a value already registered for this operator is an error.

pub fn replace_attr<T: Into<Any>>( self, key: AttrKey<T>, value: impl Into<T>, ) -> Self

Register the attribute key, replacing a value already registered for this operator (C++ set_attr(key, value, true)).

pub fn reset_attr<T>(self, key: AttrKey<T>) -> Self

Remove this operator’s value of the attribute key.

pub fn set_infer_type<F>(self, infer: F) -> Self

where F: Fn(&CallObj) -> Result<Type> + 'static,

FInferType: the result type of a Call that does not give one, which the Call constructor and Call::reinfer_type ask for; op_attr::FIXED_RETURN_TYPE takes precedence. infer is a function or a closure that captures nothing.

pub fn set_validator<F>(self, validate: F) -> Self

where F: Fn(&CallObj) -> Result<()> + 'static,

A custom validator, run by Call::validate, which reports an invalid Call as an error. It replaces the one a signature generated, and a later signature keeps it; a second custom validator is an error. validate is a function or a closure that captures nothing.

pub fn set_region_body_params<F>(self, body_params: F) -> Self

where F: Fn(&CallObj) -> Result<Array<Var>> + 'static,

FRegionGetBodyParams, which makes the operator a region operation: the variables a region binds in its body, from its Call (the operation, operands and attributes), which the RegionStmt constructor checks the body parameters against. A region that binds none returns none. body_params is a function or a closure that captures nothing.

pub fn set_region_validate<F>(self, validate: F) -> Self

where F: Fn(&RegionStmtObj) -> Result<()> + 'static,

FRegionValidate: a check of a complete region statement, run by the RegionStmt constructor after the body parameters. validate is a function or a closure that captures nothing.

Trait Implementations

impl Drop for OpDef

fn drop(&mut self)

Executes the destructor for this type.