| 199 | |
| 200 | template <typename OutputTensorType, typename InputTensorType> |
| 201 | void sparse2sparse_impl(OutputTensorType &o, const InputTensorType &a, |
| 202 | const cudaExecutor &exec) { |
| 203 | MATX_NVTX_START("", matx::MATX_NVTX_LOG_API) |
| 204 | const auto stream = exec.getStream(); |
| 205 | |
| 206 | using atype = InputTensorType; |
| 207 | using otype = OutputTensorType; |
| 208 | |
| 209 | using TA = typename atype::value_type; |
| 210 | using TO = typename otype::value_type; |
| 211 | |
| 212 | static constexpr int RANKA = atype::Rank(); |
| 213 | static constexpr int RANKO = otype::Rank(); |
| 214 | |
| 215 | // Restrictions. |
| 216 | static_assert(RANKA == 2 && RANKO == 2, "tensors must have rank-2"); |
| 217 | static_assert(std::is_same_v<TA, TO>, "tensors must have the same data type"); |
| 218 | static_assert(std::is_same_v<typename atype::crd_type, int32_t> && |
| 219 | std::is_same_v<typename otype::pos_type, int32_t> && |
| 220 | std::is_same_v<typename otype::crd_type, int32_t>, |
| 221 | "unsupported index type"); |
| 222 | |
| 223 | // Get parameters required by these tensors (for caching). |
| 224 | auto params = |
| 225 | detail::Sparse2SparseHandle_t<otype, atype>::GetConvParams(o, a, stream); |
| 226 | |
| 227 | // Lookup and cache. |
| 228 | using cache_val_type = detail::Sparse2SparseHandle_t<otype, atype>; |
| 229 | auto cache_id = detail::GetCacheIdFromType<detail::sparse2sparse_cache_t>(); |
| 230 | MATX_LOG_DEBUG("Sparse2Sparse transform: cache_id={}", cache_id); |
| 231 | detail::GetCache().LookupAndExec<detail::sparse2sparse_cache_t>( |
| 232 | cache_id, params, |
| 233 | [&]() { return std::make_shared<cache_val_type>(o, a, stream); }, |
| 234 | [&](std::shared_ptr<cache_val_type> cache_type) { |
| 235 | cache_type->Exec(o, a); |
| 236 | }, |
| 237 | exec); |
| 238 | } |
| 239 | |
| 240 | } // end namespace matx |