MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT / gemmSearch

Function gemmSearch

plugin/fcPlugin/fcPlugin.h:389–414  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

387
388template <typename T>
389cublasLtMatmulAlgo_t gemmSearch(
390 int32_t const m, int32_t const n, int32_t const k, size_t const workspaceSize, size_t& actualWorkspace)
391{
392 Gemm<T> g(m, n, k, false, false);
393 std::vector<customMatmulPerf_t> perfResults(kNB_ALGO_COMBINATIONS);
394
395 PLUGIN_CUASSERT(cudaMalloc(reinterpret_cast<void**>(&g.A), g.bytesA));
396 PLUGIN_CUASSERT(cudaMalloc(reinterpret_cast<void**>(&g.B), g.bytesB));
397 PLUGIN_CUASSERT(cudaMalloc(reinterpret_cast<void**>(&g.C), g.bytesC));
398
399 void* workspace;
400 PLUGIN_CUASSERT(cudaMalloc(&workspace, workspaceSize));
401 cublasLtHandle_t lt;
402 PLUGIN_CUBLASASSERT(cublasLtCreate(&lt));
403 LtGemmSearch(lt, g, workspace, workspaceSize, perfResults);
404 PLUGIN_CUASSERT(cudaDeviceSynchronize());
405 PLUGIN_CUBLASASSERT(cublasLtDestroy(lt));
406 PLUGIN_CUASSERT(cudaFree(workspace));
407
408 PLUGIN_CUASSERT(cudaFree(g.A));
409 PLUGIN_CUASSERT(cudaFree(g.B));
410 PLUGIN_CUASSERT(cudaFree(g.C));
411
412 actualWorkspace = perfResults[0].workspaceSize;
413 return perfResults[0].algo;
414}
415
416template <typename T>
417cublasLtMatmulAlgo_t gemmSearch(Gemm<T>& g, size_t const workspaceSize, size_t& actualWorkspace)

Callers

nothing calls this directly

Calls 1

LtGemmSearchFunction · 0.70

Tested by

no test coverage detected