larroy commented on a change in pull request #14963: Add matrix inversion
operator in linalg
URL: https://github.com/apache/incubator-mxnet/pull/14963#discussion_r284505778
##########
File path: src/operator/linalg_impl.h
##########
@@ -1233,4 +1234,252 @@ LINALG_GPU_SYEVD_WORKSPACE_QUERY(DnDsyevd, double)
#endif // __CUDACC__
+//////////////////////////////// GETRF
////////////////////////////////////////////
+
+// CPU/GPU-versions of LAPACK function "getrf"
+
+// The input of this function should be col-major for performance.
+// Tensor work holds space for ipiv in getrf
+#define LINALG_CPU_GETRF(fname, DType) \
+template<> inline \
+void linalg_getrf<cpu, DType>(const Tensor<cpu, 2, DType>& A, \
+ const Tensor<cpu, 1, DType>& work, \
+ Stream<cpu> *s) { \
+ int *ipiv = reinterpret_cast<int *>(work.dptr_); \
+ int ret(MXNET_LAPACK_##fname(MXNET_LAPACK_COL_MAJOR, A.size(1), A.size(0), \
+ A.dptr_, A.stride_, ipiv)); \
+ CHECK_EQ(ret, 0) << #fname << " failed in lapack on cpu."; \
+}
+
+LINALG_CPU_GETRF(sgetrf, float)
+LINALG_CPU_GETRF(dgetrf, double)
+
+#ifdef __CUDACC__
+
+struct set_matrix : public mxnet::op::mxnet_op::tunable {
+ template<typename DType>
+ MSHADOW_XINLINE static void Map(int i, DType **p, DType *m, int step) {
+ p[i] = m + i * step;
+ }
+};
+
+// GETRF only available with cuda8 or higher.
+#if CUDA_VERSION >= 8000
+
+// Since there is no "getri" in cuSolver, we are using batched version of
+// "getrf" and "getri" in cuBLAS here. These routines are good for large
+// batches of small matrices, so performance issue may happen when computing
+// large matices. We leave it here until MAGMA which has "getri" is introduced
+// into MXNet.
+#define LINALG_GPU_BATCH_GETRF(fname, DType) \
+template<> inline \
+void linalg_batch_getrf<gpu, DType>(const Tensor<gpu, 3, DType>& A, \
+ const Tensor<gpu, 1, DType>& work, \
+ Stream<gpu> *s) { \
+ using namespace mxnet; \
+ using namespace mxnet::op::mxnet_op; \
+ CHECK_NOTNULL(s); \
+ Storage::Handle A_ptr_buf = Storage::Get()->Alloc(sizeof(DType *) *
A.size(0), Context::GPU()); \
+ DType **A_ptr = static_cast<DType **>(A_ptr_buf.dptr); \
+ const Tensor<gpu, 3, DType> temp(work.dptr_, A.shape_, s); \
+ int *pivot = reinterpret_cast<int *>(temp.dptr_ + temp.shape_.Size()); \
+ int *info = pivot + A.size(0) * A.size(1); \
Review comment:
what happens on int32 overflow?
----------------------------------------------------------------
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:
us...@infra.apache.org
With regards,
Apache Git Services
