Searched refs:shared_memory_size (Results 1 – 5 of 5) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/core/kernels/sparse/ |
| H A D | kernels_gpu.cu.cc | 248 const size_t shared_memory_size = sizeof(int) * (batch_size + 1); in operator ()() local
251 config.thread_per_block, shared_memory_size, d.stream(), in operator ()()
316 const size_t shared_memory_size = in CSRSparseMatrixBatchMulVecImpl() local
321 config.thread_per_block, shared_memory_size, d.stream(), a_values.data(), in CSRSparseMatrixBatchMulVecImpl()
454 const size_t shared_memory_size = sizeof(int) * (batch_size + 1); in CSRSparseMatrixSoftmaxGPUImpl() local
457 shared_memory_size, d.stream(), size, rows, in CSRSparseMatrixSoftmaxGPUImpl()
659 const size_t shared_memory_size = 2 * sizeof(int) * (batch_size + 1); in CSRSparseMatrixSoftmaxGradGPUImpl() local
662 config.thread_per_block, shared_memory_size, d.stream(), size, rows, in CSRSparseMatrixSoftmaxGradGPUImpl()
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/kernels/ |
| H A D | bias_op_gpu.cu.cc | 229 int32 shared_memory_size = 0; in compute() local
231 shared_memory_size = bias_size * sizeof(typename AccumulatorType<T>::type); in compute()
234 if (shared_memory_size <= max_shared_memory_size) { in compute()
238 shared_memory_size, d.stream(), total_count, in compute()
|
| H A D | topk_op_gpu.h | 385 constexpr auto shared_memory_size = 48 << 10; // 48 KB in LaunchTopKKernel() local
388 num_shards = shared_memory_size / heap_size - 1; in LaunchTopKKernel()
412 auto shared_memory_size = (num_shards + 1) * k * sizeof(Entry<T>); variable
415 shared_memory_size, stream, input, length, k,
|
| H A D | depthwise_conv_op_gpu.h | 662 const int shared_memory_size =
667 num_outputs, device, kernel, shared_memory_size,
670 shared_memory_size, device.stream(), args, input,
1553 const int shared_memory_size =
1555 if (shared_memory_size > device.sharedMemPerBlock()) {
1583 num_out_backprop, device, kernel, shared_memory_size,
1586 shared_memory_size, device.stream(), args,
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/ |
| H A D | Reduce.cuh | 186 int shared_memory_size() const { in shared_memory_size() function
885 int shared_memory = config.shared_memory_size(); in launch_reduce_kernel()
910 int shared_memory = config.shared_memory_size(); in launch_jitted_reduce_kernel()
|