From 38b8949e1cb11ea9295955299abf6ce3e5c82acd Mon Sep 17 00:00:00 2001 From: Dr-Noob Date: Tue, 30 Nov 2021 16:03:36 +0100 Subject: [PATCH] [v0.11] Fix tensor cores calculation for Ampere. Add a brief explanation --- src/cuda/cuda.cpp | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/cuda/cuda.cpp b/src/cuda/cuda.cpp index c75efb9..21949b4 100644 --- a/src/cuda/cuda.cpp +++ b/src/cuda/cuda.cpp @@ -109,8 +109,12 @@ int64_t get_peak_performance(struct gpu_info* gpu) { } // Compute peak performance when using tensor cores -int64_t get_peak_performance_t(struct gpu_info* gpu) { - return gpu->freq * 1000000 * 4 * 4 * 8 * gpu->topo->tensor_cores; +int64_t get_peak_performance_t(cudaDeviceProp prop, struct gpu_info* gpu) { + // Volta / Turing tensor cores performs 4x4x4 FP16 matrix multiplication + // Ampere tensor cores performs 8x4x8 FP16 matrix multiplicacion + if(prop.major == 7) return gpu->freq * 1000000 * 4 * 4 * 4 * 2 * gpu->topo->tensor_cores; + else if(prop.major == 8) return gpu->freq * 1000000 * 8 * 4 * 8 * 2 * gpu->topo->tensor_cores; + else return 0; } struct gpu_info* get_gpu_info(int gpu_idx) { @@ -162,7 +166,7 @@ struct gpu_info* get_gpu_info(int gpu_idx) { gpu->mem = get_memory_info(gpu, deviceProp); gpu->topo = get_topology_info(deviceProp); gpu->peak_performance = get_peak_performance(gpu); - gpu->peak_performance_t = get_peak_performance_t(gpu); + gpu->peak_performance_t = get_peak_performance_t(deviceProp, gpu); return gpu; }