| 53 | } |
| 54 | |
| 55 | WorkspaceBundle PoolingForwardImpl::AlgoBase::get_workspace_bundle( |
| 56 | void* ptr, const SizeArgs& args) const { |
| 57 | SmallVector<size_t> sizes; |
| 58 | TensorLayout fsrc = *args.layout_src; |
| 59 | TensorLayout fdst = *args.layout_dst; |
| 60 | auto get_workspace = [&sizes](TensorLayout& layout) { |
| 61 | if (layout.dtype == dtype::BFloat16()) { |
| 62 | layout.dtype = dtype::Float32(); |
| 63 | sizes.push_back(layout.span().dist_byte()); |
| 64 | } |
| 65 | }; |
| 66 | get_workspace(fsrc); |
| 67 | get_workspace(fdst); |
| 68 | return {ptr, std::move(sizes)}; |
| 69 | } |
| 70 | |
| 71 | size_t PoolingForwardImpl::AlgoBase::get_workspace_in_bytes( |
| 72 | const SizeArgs& args) const { |
nothing calls this directly
no test coverage detected