Skip to content

batch_matmul

Classes

BatchMatMulOperator 

BatchMatMulOperator(op: AirOperator, model: AirModel, platform: SocPlatform, prefix: str = 'aot', attributes: dict[str, str] = {})

BATCH_MATMUL operator.

This operator computes the batch matrix multiplication of two tensors.

Parameters:

  • op

    (AirOperator) –

    The AIR operator to wrap.

  • model

    (AirModel) –

    The AIR model.

  • platform

    (SocPlatform) –

    The target platform for code generation.

  • prefix

    (str, default: 'aot' ) –

    Prefix for generated code files. Defaults to "aot".

  • attributes

    (dict[str, str], default: {} ) –

    Attributes for template values. Defaults to {}.

Functions

ctx_buf_required 
ctx_buf_required() -> bool

Whether the dispatched batch-matmul kernel dereferences ctx->buf.

Stated directly against the upstream preprocessor gate rather than against :meth:compute_scratch_size, so the resolve-time guard is a real cross-check of two independent derivations rather than a tautology: in arm_batch_matmul_s8.c (v7.29.0) the vector-sum buffer read lives entirely inside #if defined(ARM_MATH_MVEI).

Returns:

  • bool –

    True for an int8 output on an MVE target.

compute_scratch_size 
compute_scratch_size() -> int

Compute the scratch size for the operator.

Only arm_batch_matmul_s8 consumes a ctx buffer, and only under MVE: it needs a kernel-sum buffer of rhs_dims.w int32 words, which arm_vector_sum_s8 clears in full. The LHS never enters the requirement. arm_batch_matmul_s16 discards ctx outright ((void)ctx;) and the float kernels take no scratch, so every non-int8 dtype must report 0 rather than reserve dead arena bytes.

Returns:

  • int ( int ) –

    Computed scratch size in bytes.

validate 
validate()

Validate operand dtypes and ranks against the emitted kernel contract.

Raises:

  • ConfigValueError –

    If a tensor dtype is unsupported, the operand dtypes disagree, or an operand rank falls outside 2..4.

compute_values 
compute_values() -> dict[str, str]

Compute the values for the operator.

Returns:

  • dict[str, str] –

    dict[str, str]: Computed values for the operator.

emit 
emit(save_path: Path)

Generate the source code for the operator.

Parameters:

  • save_path
    (Path) –

    Path to save the generated code.

Functions