diff --git a/gpt-utils/gpt-utils.cpp b/gpt-utils/gpt-utils.cpp
index 1e8c9a2..58fc93a 100644
--- a/gpt-utils/gpt-utils.cpp
+++ b/gpt-utils/gpt-utils.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013,2016, The Linux Foundation. All rights reserved.
+ * Copyright (c) 2013,2016,2020 The Linux Foundation. All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are
@@ -37,10 +37,6 @@
 #include <errno.h>
 #include <sys/stat.h>
 #include <sys/ioctl.h>
-#ifndef _GENERIC_KERNEL_HEADERS
-#include <scsi/ufs/ioctl.h>
-#include <scsi/ufs/ufs.h>
-#endif
 #include <unistd.h>
 #include <linux/fs.h>
 #include <limits.h>
@@ -76,8 +72,6 @@
 #define XBL_AB_SECONDARY    "/dev/block/bootdevice/by-name/xbl_b"
 /* GPT defines */
 #define MAX_LUNS                    26
-//Size of the buffer that needs to be passed to the UFS ioctl
-#define UFS_ATTR_DATA_SIZE          32
 //This will allow us to get the root lun path from the path to the partition.
 //i.e: from /dev/block/sdaXXX get /dev/block/sda. The assumption here is that
 //the boot critical luns lie between sda to sdz which is acceptable because
@@ -129,6 +123,7 @@
      uint32_t num_valid_entries;
 };
 
+int32_t set_boot_lun(char *sg_dev,uint8_t boot_lun_id);
 /******************************************************************************
  * FUNCTIONS
  ******************************************************************************/
@@ -612,53 +607,7 @@
         return -1;
 }
 
-int set_boot_lun(char *sg_dev, uint8_t boot_lun_id)
-{
-#ifndef _GENERIC_KERNEL_HEADERS
-        int fd = -1;
-        int rc;
-        struct ufs_ioctl_query_data *data = NULL;
-        size_t ioctl_data_size = sizeof(struct ufs_ioctl_query_data) + UFS_ATTR_DATA_SIZE;
 
-        data = (struct ufs_ioctl_query_data*)malloc(ioctl_data_size);
-        if (!data) {
-                fprintf(stderr, "%s: Failed to alloc query data struct\n",
-                                __func__);
-                goto error;
-        }
-        memset(data, 0, ioctl_data_size);
-        data->opcode = UPIU_QUERY_OPCODE_WRITE_ATTR;
-        data->idn = QUERY_ATTR_IDN_BOOT_LU_EN;
-        data->buf_size = UFS_ATTR_DATA_SIZE;
-        data->buffer[0] = boot_lun_id;
-        fd = open(sg_dev, O_RDWR);
-        if (fd < 0) {
-                fprintf(stderr, "%s: Failed to open %s(%s)\n",
-                                __func__,
-                                sg_dev,
-                                strerror(errno));
-                goto error;
-        }
-        rc = ioctl(fd, UFS_IOCTL_QUERY, data);
-        if (rc) {
-                fprintf(stderr, "%s: UFS query ioctl failed(%s)\n",
-                                __func__,
-                                strerror(errno));
-                goto error;
-        }
-        close(fd);
-        free(data);
-        return 0;
-error:
-        if (fd >= 0)
-                close(fd);
-        if (data)
-                free(data);
-        return -1;
-#else
-	return 0;
-#endif
-}
 
 //Swtich betwieen using either the primary or the backup
 //boot LUN for boot. This is required since UFS boot partitions
@@ -732,6 +681,7 @@
                                 __func__);
                 goto error;
         }
+        /* set boot lun using /dev/sg or /dev/ufs-bsg* */
         if (set_boot_lun(sg_dev_node, boot_lun_id)) {
                 fprintf(stderr, "%s: Failed to set xblbak as boot partition\n",
                                 __func__);