Home
last modified time | relevance | path

Searched refs:fbgemm_kernel_cache (Results 1 – 2 of 2) sorted by relevance

/aosp_15_r20/external/pytorch/aten/src/ATen/native/
H A DEmbeddingBag.cpp200 _EmbeddingBagKernelCache* fbgemm_kernel_cache) { in index_select_add() argument
228 auto kernel_16bit_index_t = fbgemm_kernel_cache in index_select_add()
229 ? fbgemm_kernel_cache in index_select_add()
378 _EmbeddingBagKernelCache* fbgemm_kernel_cache) { in index_select_add() argument
407 fbgemm_kernel_cache ? in index_select_add()
408 fbgemm_kernel_cache->getCallback</* has_weight */ false, index_t, float>(ddim) : in index_select_add()
565 _EmbeddingBagKernelCache* fbgemm_kernel_cache) { in index_select_scale_add() argument
607 auto kernel_16bit_index_t = fbgemm_kernel_cache in index_select_scale_add()
608 ? fbgemm_kernel_cache in index_select_scale_add()
754 _EmbeddingBagKernelCache* fbgemm_kernel_cache) { in index_select_scale_add() argument
[all …]
H A DEmbeddingBag.h135 _EmbeddingBagKernelCache* fbgemm_kernel_cache = nullptr);
151 _EmbeddingBagKernelCache* fbgemm_kernel_cache = nullptr);