diff --git a/drivers/scsi/mpt3sas/mpt3sas_ctl.c b/drivers/scsi/mpt3sas/mpt3sas_ctl.c
index 9b37a32..a14ff88 100644
--- a/drivers/scsi/mpt3sas/mpt3sas_ctl.c
+++ b/drivers/scsi/mpt3sas/mpt3sas_ctl.c
@@ -1491,6 +1491,26 @@ _ctl_diag_capability(struct MPT3SAS_ADAPTER *ioc, u8 buffer_type)
 	return rc;
 }
 
+/**
+ * _ctl_diag_get_bufftype - return diag buffer type
+ *              either TRACE, SNAPSHOT, or EXTENDED
+ * @ioc: per adapter object
+ * @unique_id: specifies the unique_id for the buffer
+ *
+ * returns MPT3_DIAG_UID_NOT_FOUND if the id not found
+ */
+static u8
+_ctl_diag_get_bufftype(struct MPT3SAS_ADAPTER *ioc, u32 unique_id)
+{
+	u8  index;
+
+	for (index = 0; index < MPI2_DIAG_BUF_TYPE_COUNT; index++) {
+		if (ioc->unique_id[index] == unique_id)
+			return index;
+	}
+
+	return MPT3_DIAG_UID_NOT_FOUND;
+}
 
 /**
  * _ctl_diag_register_2 - wrapper for registering diag buffer support
@@ -1538,11 +1558,65 @@ _ctl_diag_register_2(struct MPT3SAS_ADAPTER *ioc,
 		return -EPERM;
 	}
 
+	if (diag_register->unique_id == 0) {
+		ioc_err(ioc,
+		    "%s: Invalid UID(0x%08x), buffer_type(0x%02x)\n", __func__,
+		    diag_register->unique_id, buffer_type);
+		return -EINVAL;
+	}
+
 	if (ioc->diag_buffer_status[buffer_type] &
 	    MPT3_DIAG_BUFFER_IS_REGISTERED) {
-		ioc_err(ioc, "%s: already has a registered buffer for buffer_type(0x%02x)\n",
-			__func__, buffer_type);
-		return -EINVAL;
+		/*
+		 * If driver posts buffer initially, then an application wants
+		 * to Register that buffer (own it) without Releasing first,
+		 * the application Register command MUST have the same buffer
+		 * type and size in the Register command (obtained from the
+		 * Query command). Otherwise that Register command will be
+		 * failed. If the application has released the buffer but wants
+		 * to re-register it, it should be allowed as long as the
+		 * Unique-Id/Size match.
+		 */
+
+		if (ioc->unique_id[buffer_type] == MPT3DIAGBUFFUNIQUEID &&
+		    ioc->diag_buffer_sz[buffer_type] ==
+		    diag_register->requested_buffer_size) {
+
+			if (!(ioc->diag_buffer_status[buffer_type] &
+			     MPT3_DIAG_BUFFER_IS_RELEASED)) {
+				dctlprintk(ioc, ioc_info(ioc,
+				    "%s: diag_buffer (%d) ownership changed. old-ID(0x%08x), new-ID(0x%08x)\n",
+				    __func__, buffer_type,
+				    ioc->unique_id[buffer_type],
+				    diag_register->unique_id));
+
+				/*
+				 * Application wants to own the buffer with
+				 * the same size.
+				 */
+				ioc->unique_id[buffer_type] =
+				    diag_register->unique_id;
+				rc = 0; /* success */
+				goto out;
+			}
+		} else if (ioc->unique_id[buffer_type] !=
+		    MPT3DIAGBUFFUNIQUEID) {
+			if (ioc->unique_id[buffer_type] !=
+			    diag_register->unique_id ||
+			    ioc->diag_buffer_sz[buffer_type] !=
+			    diag_register->requested_buffer_size ||
+			    !(ioc->diag_buffer_status[buffer_type] &
+			    MPT3_DIAG_BUFFER_IS_RELEASED)) {
+				ioc_err(ioc,
+				    "%s: already has a registered buffer for buffer_type(0x%02x)\n",
+				    __func__, buffer_type);
+				return -EINVAL;
+			}
+		} else {
+			ioc_err(ioc, "%s: already has a registered buffer for buffer_type(0x%02x)\n",
+			    __func__, buffer_type);
+			return -EINVAL;
+		}
 	}
 
 	if (diag_register->requested_buffer_size % 4)  {
@@ -1689,7 +1763,9 @@ mpt3sas_enable_diag_buffer(struct MPT3SAS_ADAPTER *ioc, u8 bits_to_register)
 		ioc->diag_trigger_master.MasterData =
 		    (MASTER_TRIGGER_FW_FAULT + MASTER_TRIGGER_ADAPTER_RESET);
 		diag_register.buffer_type = MPI2_DIAG_BUF_TYPE_TRACE;
-		diag_register.unique_id = 0x7075900;
+		diag_register.unique_id =
+		    (ioc->hba_mpi_version_belonged == MPI2_VERSION) ?
+		    (MPT2DIAGBUFFUNIQUEID):(MPT3DIAGBUFFUNIQUEID);
 
 		if (trace_buff_size != 0) {
 			diag_register.requested_buffer_size = trace_buff_size;
@@ -1815,7 +1891,13 @@ _ctl_diag_unregister(struct MPT3SAS_ADAPTER *ioc, void __user *arg)
 	dctlprintk(ioc, ioc_info(ioc, "%s\n",
 				 __func__));
 
-	buffer_type = karg.unique_id & 0x000000ff;
+	buffer_type = _ctl_diag_get_bufftype(ioc, karg.unique_id);
+	if (buffer_type == MPT3_DIAG_UID_NOT_FOUND) {
+		ioc_err(ioc, "%s: buffer with unique_id(0x%08x) not found\n",
+		    __func__, karg.unique_id);
+		return -EINVAL;
+	}
+
 	if (!_ctl_diag_capability(ioc, buffer_type)) {
 		ioc_err(ioc, "%s: doesn't have capability for buffer_type(0x%02x)\n",
 			__func__, buffer_type);
@@ -1899,7 +1981,7 @@ _ctl_diag_query(struct MPT3SAS_ADAPTER *ioc, void __user *arg)
 		return -EINVAL;
 	}
 
-	if (karg.unique_id & 0xffffff00) {
+	if (karg.unique_id) {
 		if (karg.unique_id != ioc->unique_id[buffer_type]) {
 			ioc_err(ioc, "%s: unique_id(0x%08x) is not registered\n",
 				__func__, karg.unique_id);
@@ -2065,7 +2147,13 @@ _ctl_diag_release(struct MPT3SAS_ADAPTER *ioc, void __user *arg)
 	dctlprintk(ioc, ioc_info(ioc, "%s\n",
 				 __func__));
 
-	buffer_type = karg.unique_id & 0x000000ff;
+	buffer_type = _ctl_diag_get_bufftype(ioc, karg.unique_id);
+	if (buffer_type == MPT3_DIAG_UID_NOT_FOUND) {
+		ioc_err(ioc, "%s: buffer with unique_id(0x%08x) not found\n",
+		    __func__, karg.unique_id);
+		return -EINVAL;
+	}
+
 	if (!_ctl_diag_capability(ioc, buffer_type)) {
 		ioc_err(ioc, "%s: doesn't have capability for buffer_type(0x%02x)\n",
 			__func__, buffer_type);
@@ -2149,7 +2237,13 @@ _ctl_diag_read_buffer(struct MPT3SAS_ADAPTER *ioc, void __user *arg)
 	dctlprintk(ioc, ioc_info(ioc, "%s\n",
 				 __func__));
 
-	buffer_type = karg.unique_id & 0x000000ff;
+	buffer_type = _ctl_diag_get_bufftype(ioc, karg.unique_id);
+	if (buffer_type == MPT3_DIAG_UID_NOT_FOUND) {
+		ioc_err(ioc, "%s: buffer with unique_id(0x%08x) not found\n",
+		    __func__, karg.unique_id);
+		return -EINVAL;
+	}
+
 	if (!_ctl_diag_capability(ioc, buffer_type)) {
 		ioc_err(ioc, "%s: doesn't have capability for buffer_type(0x%02x)\n",
 			__func__, buffer_type);
@@ -3202,7 +3296,9 @@ host_trace_buffer_enable_store(struct device *cdev,
 		} else
 			diag_register.requested_buffer_size = (1024 * 1024);
 
-		diag_register.unique_id = 0x7075900;
+		diag_register.unique_id =
+		    (ioc->hba_mpi_version_belonged == MPI2_VERSION) ?
+		    (MPT2DIAGBUFFUNIQUEID):(MPT3DIAGBUFFUNIQUEID);
 		ioc->diag_buffer_status[MPI2_DIAG_BUF_TYPE_TRACE] = 0;
 		_ctl_diag_register_2(ioc,  &diag_register);
 	} else if (!strcmp(str, "release")) {
diff --git a/drivers/scsi/mpt3sas/mpt3sas_ctl.h b/drivers/scsi/mpt3sas/mpt3sas_ctl.h
index 18b46fa..d1a6ab1 100644
--- a/drivers/scsi/mpt3sas/mpt3sas_ctl.h
+++ b/drivers/scsi/mpt3sas/mpt3sas_ctl.h
@@ -95,6 +95,14 @@
 #define MPT3DIAGREADBUFFER _IOWR(MPT3_MAGIC_NUMBER, 30, \
 	struct mpt3_diag_read_buffer)
 
+/* Trace Buffer default UniqueId */
+#define MPT2DIAGBUFFUNIQUEID (0x07075900)
+#define MPT3DIAGBUFFUNIQUEID (0x4252434D)
+
+/* UID not found */
+#define MPT3_DIAG_UID_NOT_FOUND (0xFF)
+
+
 /**
  * struct mpt3_ioctl_header - main header structure
  * @ioc_number -  IOC unit number
