Method run

src/cpu/operators/CpuMatMul.cpp:288–336

}

void CpuMatMul::run(ITensorPack &tensors)
{
    ARM_COMPUTE_TRACE_EVENT(ARM_COMPUTE_PROF_CAT_CPU, ARM_COMPUTE_PROF_LVL_CPU, "CpuMatMul::run");
    // Retrieve tensors from tensor pack
    auto lhs = tensors.get_tensor(ACL_SRC_0);
    auto rhs = tensors.get_const_tensor(ACL_SRC_1);
    auto dst = tensors.get_tensor(ACL_DST);

    // Reshape LHS and DST to ensure compatibility with GEMM asm kernel (Batch dimensions is 4th for lhs and dst within asm)
    // Collapse RHS (necessary to support dimensions larger than 3 in gemm assembly)
    lhs->info()->set_tensor_shape(
        TensorShape(_original_lhs_shape.x(), _original_lhs_shape.y(), 1,
                    _original_lhs_shape.collapsed_from(2).z())); // Collapsed 3+ dimensions into z
    dst->info()->set_tensor_shape(
        TensorShape(_original_dst_shape.x(), _original_dst_shape.y(), 1,
                    _original_dst_shape.collapsed_from(2).z())); // Collapsed 3+ dimensions into z
    rhs->info()->set_tensor_shape(_original_rhs_shape.collapsed_from(2));

    // Initialise object to handle stored transposed tensors in auxillary memory
    CpuAuxTensorHandler lhs_transposed(offset_int_vec(TransposeLHS), _lhs_transposed, tensors, true);
    CpuAuxTensorHandler rhs_transposed(offset_int_vec(TransposeRHS), _rhs_transposed, tensors, true);

    // Create tensor pack for asm kernel
    ITensorPack asm_tensors(tensors);

    // Run transpose lhs if necessary
    if (_adj_lhs)
    {
        ITensorPack lhs_transpose_pack = {{TensorType::ACL_SRC, lhs}, {TensorType::ACL_DST, lhs_transposed.get()}};
        NEScheduler::get().schedule_op(_transpose_kernel_lhs.get(), Window::DimY, _transpose_kernel_lhs->window(),
                                       lhs_transpose_pack);
        asm_tensors.add_const_tensor(TensorType::ACL_SRC_0, lhs_transposed.get());
    }
    // Run transpose rhs if necessary
    if (_adj_rhs)
    {
        ITensorPack rhs_transpose_pack = {{TensorType::ACL_SRC, rhs}, {TensorType::ACL_DST, rhs_transposed.get()}};
        NEScheduler::get().schedule_op(_transpose_kernel_rhs.get(), Window::DimY, _transpose_kernel_rhs->window(),
                                       rhs_transpose_pack);
        asm_tensors.add_const_tensor(TensorType::ACL_SRC_1, rhs_transposed.get());
    }
    // Run asm kernel
    _asm_glue->run(asm_tensors);

    // Undo reshape of tensors
    dst->info()->set_tensor_shape(_original_dst_shape);
    lhs->info()->set_tensor_shape(_original_lhs_shape);
    rhs->info()->set_tensor_shape(_original_rhs_shape);
}

experimental::MemoryRequirements CpuMatMul::workspace() const
{