[mlir][gpu][sparse] gracefully accept zero size allocation #66127
+10
−17
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This cleans up a unnecessary code that changes zero size allocation to avoid the following error message 'cuMemAlloc(&ptr, sizeBytes)' failed with 'CUDA_ERROR_INVALID_VALUE'
|
@llvm/pr-subscribers-mlir-execution-engine ChangesThis cleans up a unnecessary code that changes zero size allocation to avoid the following error message 'cuMemAlloc(&ptr, sizeBytes)' failed with 'CUDA_ERROR_INVALID_VALUE'Full diff: https://github.com/llvm/llvm-project/pull/66127.diff 1 Files Affected:
diff --git a/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp b/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp
index 1dba677ebe66365..7bf6804902479a8 100644
--- a/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp
+++ b/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp
@@ -212,8 +212,9 @@ extern MLIR_CUDA_WRAPPERS_EXPORT "C" void mgpuEventRecord(CUevent event,
extern "C" void *mgpuMemAlloc(uint64_t sizeBytes, CUstream /*stream*/) {
ScopedContext scopedContext;
- CUdeviceptr ptr;
- CUDA_REPORT_IF_ERROR(cuMemAlloc(&ptr, sizeBytes));
+ CUdeviceptr ptr = 0;
+ if (sizeBytes != 0)
+ CUDA_REPORT_IF_ERROR(cuMemAlloc(&ptr, sizeBytes));
return reinterpret_cast(ptr);
}
@@ -521,7 +522,7 @@ extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t mgpuSpMVBufferSize(
CUSPARSE_REPORT_IF_ERROR(cusparseSpMV_bufferSize(
cusparse_env, modeA, alphap, matA, vecX, betap, vecY, cTp,
CUSPARSE_SPMV_ALG_DEFAULT, &bufferSize))
- return bufferSize == 0 ? 1 : bufferSize; // avoid zero-alloc
+ return bufferSize;
}
extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSpMV(int32_t ma, void *a, void *x,
@@ -555,7 +556,7 @@ mgpuSpMMBufferSize(int32_t ma, int32_t mb, void *a, void *b, void *c,
CUSPARSE_REPORT_IF_ERROR(cusparseSpMM_bufferSize(
cusparse_env, modeA, modeB, alphap, matA, matB, betap, matC, cTp,
CUSPARSE_SPMM_ALG_DEFAULT, &bufferSize))
- return bufferSize == 0 ? 1 : bufferSize; // avoid zero-alloc
+ return bufferSize;
}
extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSpMM(int32_t ma, int32_t mb,
@@ -590,7 +591,7 @@ mgpuSDDMMBufferSize(int32_t ma, int32_t mb, void *a, void *b, void *c,
CUSPARSE_REPORT_IF_ERROR(cusparseSDDMM_bufferSize(
cusparse_env, modeA, modeB, alphap, matA, matB, betap, matC, cTp,
CUSPARSE_SDDMM_ALG_DEFAULT, &bufferSize))
- return bufferSize == 0 ? 1 : bufferSize; // avoid zero-alloc
+ return bufferSize;
}
extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSDDMM(int32_t ma, int32_t mb,
@@ -638,7 +639,7 @@ extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t mgpuSpGEMMWorkEstimation(
CUSPARSE_REPORT_IF_ERROR(cusparseSpGEMM_workEstimation(
cusparse_env, modeA, modeB, alphap, matA, matB, betap, matC, cTp,
CUSPARSE_SPGEMM_DEFAULT, spgemmDesc, &newBufferSize, buf))
- return newBufferSize == 0 ? 1 : newBufferSize; // avoid zero-alloc
+ return newBufferSize;
}
extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t
@@ -656,7 +657,7 @@ mgpuSpGEMMCompute(void *s, int32_t ma, int32_t mb, void *a, void *b, void *c,
CUSPARSE_REPORT_IF_ERROR(cusparseSpGEMM_compute(
cusparse_env, modeA, modeB, alphap, matA, matB, betap, matC, cTp,
CUSPARSE_SPGEMM_DEFAULT, spgemmDesc, &newBufferSize2, buf2))
- return newBufferSize2 == 0 ? 1 : newBufferSize2; // avoid zero-alloc
+ return newBufferSize2;
}
extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
@@ -794,7 +795,6 @@ mgpuCuSparseLtSpMMBufferSize(void *bs, int32_t ma, int32_t mb, void *a, void *b,
auto workspace_size = reinterpret_cast(bs);
auto compressed_size = &(reinterpret_cast(bs)[1]);
auto compressed_buffer_size = &(reinterpret_cast(bs)[2]);
- size_t workspace_size_, compressed_size_, compressed_buffer_size_;
auto cTp = static_cast(ctp);
cusparseOperation_t modeA = static_cast(ma);
@@ -836,16 +836,9 @@ mgpuCuSparseLtSpMMBufferSize(void *bs, int32_t ma, int32_t mb, void *a, void *b,
}
CUSPARSE_REPORT_IF_ERROR(cusparseLtMatmulGetWorkspace(
- &cusparseLt_env, &(matA->plan), &workspace_size_))
+ &cusparseLt_env, &(matA->plan), workspace_size))
CUSPARSE_REPORT_IF_ERROR(cusparseLtSpMMACompressedSize(
- &cusparseLt_env, &(matA->plan), &compressed_size_,
- &compressed_buffer_size_))
-
- // Avoid zero-allocation.
- *workspace_size = (workspace_size_ == 0 ? 1 : workspace_size_);
- *compressed_size = (compressed_size_ == 0 ? 1 : compressed_size_);
- *compressed_buffer_size =
- (compressed_buffer_size_ == 0 ? 1 : compressed_buffer_size_);
+ &cusparseLt_env, &(matA->plan), compressed_size, compressed_buffer_size))
}
extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
|
|
@llvm/pr-subscribers-mlir ChangesThis cleans up a unnecessary code that changes zero size allocation to avoid the following error message 'cuMemAlloc(&ptr, sizeBytes)' failed with 'CUDA_ERROR_INVALID_VALUE'Full diff: https://github.com/llvm/llvm-project/pull/66127.diff 1 Files Affected:
diff --git a/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp b/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp
index 1dba677ebe66365..7bf6804902479a8 100644
--- a/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp
+++ b/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp
@@ -212,8 +212,9 @@ extern MLIR_CUDA_WRAPPERS_EXPORT "C" void mgpuEventRecord(CUevent event,
extern "C" void *mgpuMemAlloc(uint64_t sizeBytes, CUstream /*stream*/) {
ScopedContext scopedContext;
- CUdeviceptr ptr;
- CUDA_REPORT_IF_ERROR(cuMemAlloc(&ptr, sizeBytes));
+ CUdeviceptr ptr = 0;
+ if (sizeBytes != 0)
+ CUDA_REPORT_IF_ERROR(cuMemAlloc(&ptr, sizeBytes));
return reinterpret_cast(ptr);
}
@@ -521,7 +522,7 @@ extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t mgpuSpMVBufferSize(
CUSPARSE_REPORT_IF_ERROR(cusparseSpMV_bufferSize(
cusparse_env, modeA, alphap, matA, vecX, betap, vecY, cTp,
CUSPARSE_SPMV_ALG_DEFAULT, &bufferSize))
- return bufferSize == 0 ? 1 : bufferSize; // avoid zero-alloc
+ return bufferSize;
}
extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSpMV(int32_t ma, void *a, void *x,
@@ -555,7 +556,7 @@ mgpuSpMMBufferSize(int32_t ma, int32_t mb, void *a, void *b, void *c,
CUSPARSE_REPORT_IF_ERROR(cusparseSpMM_bufferSize(
cusparse_env, modeA, modeB, alphap, matA, matB, betap, matC, cTp,
CUSPARSE_SPMM_ALG_DEFAULT, &bufferSize))
- return bufferSize == 0 ? 1 : bufferSize; // avoid zero-alloc
+ return bufferSize;
}
extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSpMM(int32_t ma, int32_t mb,
@@ -590,7 +591,7 @@ mgpuSDDMMBufferSize(int32_t ma, int32_t mb, void *a, void *b, void *c,
CUSPARSE_REPORT_IF_ERROR(cusparseSDDMM_bufferSize(
cusparse_env, modeA, modeB, alphap, matA, matB, betap, matC, cTp,
CUSPARSE_SDDMM_ALG_DEFAULT, &bufferSize))
- return bufferSize == 0 ? 1 : bufferSize; // avoid zero-alloc
+ return bufferSize;
}
extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSDDMM(int32_t ma, int32_t mb,
@@ -638,7 +639,7 @@ extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t mgpuSpGEMMWorkEstimation(
CUSPARSE_REPORT_IF_ERROR(cusparseSpGEMM_workEstimation(
cusparse_env, modeA, modeB, alphap, matA, matB, betap, matC, cTp,
CUSPARSE_SPGEMM_DEFAULT, spgemmDesc, &newBufferSize, buf))
- return newBufferSize == 0 ? 1 : newBufferSize; // avoid zero-alloc
+ return newBufferSize;
}
extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t
@@ -656,7 +657,7 @@ mgpuSpGEMMCompute(void *s, int32_t ma, int32_t mb, void *a, void *b, void *c,
CUSPARSE_REPORT_IF_ERROR(cusparseSpGEMM_compute(
cusparse_env, modeA, modeB, alphap, matA, matB, betap, matC, cTp,
CUSPARSE_SPGEMM_DEFAULT, spgemmDesc, &newBufferSize2, buf2))
- return newBufferSize2 == 0 ? 1 : newBufferSize2; // avoid zero-alloc
+ return newBufferSize2;
}
extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
@@ -794,7 +795,6 @@ mgpuCuSparseLtSpMMBufferSize(void *bs, int32_t ma, int32_t mb, void *a, void *b,
auto workspace_size = reinterpret_cast(bs);
auto compressed_size = &(reinterpret_cast(bs)[1]);
auto compressed_buffer_size = &(reinterpret_cast(bs)[2]);
- size_t workspace_size_, compressed_size_, compressed_buffer_size_;
auto cTp = static_cast(ctp);
cusparseOperation_t modeA = static_cast(ma);
@@ -836,16 +836,9 @@ mgpuCuSparseLtSpMMBufferSize(void *bs, int32_t ma, int32_t mb, void *a, void *b,
}
CUSPARSE_REPORT_IF_ERROR(cusparseLtMatmulGetWorkspace(
- &cusparseLt_env, &(matA->plan), &workspace_size_))
+ &cusparseLt_env, &(matA->plan), workspace_size))
CUSPARSE_REPORT_IF_ERROR(cusparseLtSpMMACompressedSize(
- &cusparseLt_env, &(matA->plan), &compressed_size_,
- &compressed_buffer_size_))
-
- // Avoid zero-allocation.
- *workspace_size = (workspace_size_ == 0 ? 1 : workspace_size_);
- *compressed_size = (compressed_size_ == 0 ? 1 : compressed_size_);
- *compressed_buffer_size =
- (compressed_buffer_size_ == 0 ? 1 : compressed_buffer_size_);
+ &cusparseLt_env, &(matA->plan), compressed_size, compressed_buffer_size))
}
extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
|
grypp
approved these changes
Sep 12, 2023
ZijunZhaoCCK
pushed a commit
to ZijunZhaoCCK/llvm-project
that referenced
this pull request
Sep 19, 2023
This cleans up a unnecessary code that changes zero size allocation to avoid the following error message 'cuMemAlloc(&ptr, sizeBytes)' failed with 'CUDA_ERROR_INVALID_VALUE'
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This cleans up a unnecessary code that changes zero size allocation to avoid the following error message
'cuMemAlloc(&ptr, sizeBytes)' failed with 'CUDA_ERROR_INVALID_VALUE'