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

Function customMatmulRun

plugin/fcPlugin/fcPlugin.cpp:73–132  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

71}
72
73static cublasStatus_t customMatmulRun(cublasLtHandle_t ltHandle, // to get the capabilities (required a GPU)
74 cublasLtMatmulDesc_t operationDesc, void const* alpha, // host or device pointer
75 void const* A, cublasLtMatrixLayout_t Adesc, void const* B, cublasLtMatrixLayout_t Bdesc,
76 void const* beta, // host or device pointer
77 void const* C, cublasLtMatrixLayout_t Cdesc, void* D, cublasLtMatrixLayout_t Ddesc,
78 cublasLtMatmulAlgo_t const& algo, void* workSpace, size_t workSpaceSizeInBytes, customMatmulPerf_t& perfResults,
79 cudaStream_t stream, cudaEvent_t& startEvent, cudaEvent_t& stopEvent)
80{
81
82 cublasLtMatmulHeuristicResult_t heurResult;
83
84 // Looping over the Algo
85 cublasStatus_t algoStatus
86 = cublasLtMatmulAlgoCheck(ltHandle, operationDesc, Adesc, Bdesc, Cdesc, Ddesc, &algo, &heurResult);
87
88 if (algoStatus == CUBLAS_STATUS_SUCCESS)
89 {
90 if (heurResult.workspaceSize <= workSpaceSizeInBytes)
91 {
92 if (cudaEventRecord(startEvent, stream) != cudaSuccess)
93 {
94 return CUBLAS_STATUS_INTERNAL_ERROR;
95 }
96 for (int32_t loop = 0; loop < kNB_KERNEL_REPEATS; loop++)
97 {
98 cublasStatus_t oneRunStatus = cublasLtMatmul(ltHandle, operationDesc, alpha, // host or device pointer
99 A, Adesc, B, Bdesc, beta, // host or device pointer
100 C, Cdesc, D, Ddesc, &algo, workSpace, workSpaceSizeInBytes, stream);
101 if (oneRunStatus != CUBLAS_STATUS_SUCCESS)
102 {
103 algoStatus = oneRunStatus;
104 break;
105 }
106 }
107 if (cudaEventRecord(stopEvent, stream) != cudaSuccess)
108 {
109 return CUBLAS_STATUS_INTERNAL_ERROR;
110 }
111 if (cudaEventSynchronize(stopEvent) != cudaSuccess)
112 {
113 return CUBLAS_STATUS_INTERNAL_ERROR;
114 }
115 float time;
116 if (cudaEventElapsedTime(&time, startEvent, stopEvent) != cudaSuccess)
117 {
118 return CUBLAS_STATUS_INTERNAL_ERROR;
119 }
120 // For the moment only add successful findings
121 perfResults.algo = algo;
122 perfResults.time = time / kNB_KERNEL_REPEATS; // Average time
123 perfResults.workspaceSize = heurResult.workspaceSize;
124 perfResults.wavesCount = heurResult.wavesCount;
125 }
126 else
127 {
128 algoStatus = CUBLAS_STATUS_NOT_SUPPORTED; // Not enough workspace
129 }
130 }

Callers 1

LtGemmSearchFunction · 0.85

Calls 1

cublasLtMatmulFunction · 0.85

Tested by

no test coverage detected