Hzfengsy commented on a change in pull request #4550: [Perf] Add CublasLt
extern support for better Igemm performance
URL: https://github.com/apache/incubator-tvm/pull/4550#discussion_r360612582
##########
File path: src/runtime/contrib/cublas/cublas.cc
##########
@@ -181,6 +181,98 @@ bool CheckMixPrecisionType(DLDataType in_dtype,
DLDataType out_dtype, bool int_s
}
}
+int roundoff(int v, int d) {
+ return (v + d - 1) / d * d;
+}
+
+#if CUDART_VERSION >= 10010
+inline void CallLtIgemm(TVMArgs args, TVMRetValue *ret, cublasLtHandle_t hdl) {
+ DLTensor *A = args[0];
+ DLTensor *B = args[1];
+ DLTensor *C = args[2];
+ bool transa = args[3];
+ bool transb = args[4];
+ // Reversed strides indicates an in-place transpose operation.
+ transa = IsInPlaceTransposed(A) ? !transa : transa;
+ transb = IsInPlaceTransposed(B) ? !transb : transb;
+ int M = ColumnCount(B, transb);
+ int N = RowCount(A, transa);
+ int K = ColumnCount(A, transa);
+ int N_out = ColumnCount(C, false);
+ int m = M;
+ int n = m;
+ int k = m;
+ int lda = M * K / (roundoff(K, 32) / 32);
+ int ldb = K * N / (roundoff(K, 32) / 32);
+ int ldc = M * N_out / (roundoff(N_out, 32) / 32);
+ CHECK_EQ(A->ndim, 2);
+ CHECK_EQ(B->ndim, 2);
+ CHECK_EQ(C->ndim, 2);
+
+ CHECK_EQ(ElementStride(A), 1);
+ CHECK_EQ(ElementStride(B), 1);
+ CHECK_EQ(ElementStride(C), 1);
+
+ CHECK(TypeEqual(A->dtype, B->dtype));
+ CHECK(TypeMatch(A->dtype, kDLInt, 8));
+ CHECK(TypeMatch(C->dtype, kDLInt, 32));
+
+ CHECK(CheckMixPrecisionType(A->dtype, C->dtype)) << "Unsupported data type";
+ int32_t alpha = args.size() > 5 ? args[5] : 1;
+ int32_t beta = args.size() > 6 ? args[6] : 0;
+ cublasLtMatrixLayout_t Adesc = NULL, Bdesc = NULL, Cdesc = NULL;
+ auto A_data = reinterpret_cast<void *>(static_cast<char *>(A->data) +
A->byte_offset);
Review comment:
style
```suggestion
auto A_data = reinterpret_cast<void*>(static_cast<char*>(A->data) +
A->byte_offset);
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
