From 0ec4bd58e50c6d96098328d12ad86aa95dbb71f3 Mon Sep 17 00:00:00 2001
From: Mahesh Doijade <mdoijade@nvidia.com>
Date: Mon, 1 Jun 2020 20:22:22 +0530
Subject: [PATCH] cudaCompressibleMemory: when unable to allocate compressible
 memory waive the execution

---
 Samples/cudaCompressibleMemory/compMalloc.cpp | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/Samples/cudaCompressibleMemory/compMalloc.cpp b/Samples/cudaCompressibleMemory/compMalloc.cpp
index 7503e3d3..78f0a43d 100644
--- a/Samples/cudaCompressibleMemory/compMalloc.cpp
+++ b/Samples/cudaCompressibleMemory/compMalloc.cpp
@@ -27,6 +27,7 @@
 
 #include <stdio.h>
 #include <string.h>
+#include <helper_cuda.h>
 #include <cuda.h>
 #include <cuda_runtime_api.h>
 
@@ -67,6 +68,16 @@ cudaError_t allocateCompressible(void **adr, size_t size, bool UseCompressibleMe
     if (cuMemCreate(&allocationHandle, size, &prop, 0) != CUDA_SUCCESS)
         return cudaErrorMemoryAllocation;
 
+    // Check if cuMemCreate was able to allocate compressible memory.
+    if (UseCompressibleMemory) {
+        CUmemAllocationProp allocationProp = {};
+        cuMemGetAllocationPropertiesFromHandle(&allocationProp, allocationHandle);
+        if (allocationProp.allocFlags.compressionType != CU_MEM_ALLOCATION_COMP_GENERIC) {
+            printf("Could not allocate compressible memory... so waiving execution\n");
+            exit(EXIT_WAIVED);
+        }
+    }
+
     if (cuMemMap(dptr, size, 0, allocationHandle, 0) != CUDA_SUCCESS)
         return cudaErrorMemoryAllocation;