| 41 | } |
| 42 | |
| 43 | std::shared_ptr<Tensor<>> matmul(Handle* handle, const TensorND& A, const TensorND& B) { |
| 44 | auto matmul_opr = handle->create_operator<BatchedMatrixMul>(); |
| 45 | |
| 46 | TensorLayout result_layout; |
| 47 | matmul_opr->deduce_layout(A.layout, B.layout, result_layout); |
| 48 | std::shared_ptr<Tensor<>> result(new Tensor<>(handle, result_layout)); |
| 49 | WorkspaceWrapper ws( |
| 50 | handle, |
| 51 | matmul_opr->get_workspace_in_bytes(A.layout, B.layout, result->layout())); |
| 52 | matmul_opr->exec(A, B, result->tensornd(), ws.workspace()); |
| 53 | return result; |
| 54 | } |
| 55 | |
| 56 | } // namespace |
| 57 |
no test coverage detected