diff --git a/security/pfe/pfk.c b/security/pfe/pfk.c
index f9a6292..0cdfd30 100644
--- a/security/pfe/pfk.c
+++ b/security/pfe/pfk.c
@@ -76,7 +76,6 @@
 	struct pfk_key_info *key_info,
 	enum ice_cryto_algo_mode *algo,
 	bool *is_pfe,
-	unsigned int *data_unit,
 	const char *storage_type);
 
 typedef bool (*pfk_allow_merge_bio_type)(const struct bio *bio1,
@@ -294,13 +293,17 @@
 	which_pfe = pfk_get_pfe_type(inode);
 	s_type = (char *)pfk_kc_get_storage_type();
 
+	if (data_unit && (bio_dun(bio) ||
+			!memcmp(s_type, "ufs", strlen("ufs"))))
+		*data_unit = 1 << ICE_CRYPTO_DATA_UNIT_4_KB;
+
 	if (which_pfe != INVALID_PFE) {
 		/* Encrypted file; override ->bi_crypt_key */
 		pr_debug("parsing inode %lu with PFE type %d\n",
 			 inode->i_ino, which_pfe);
 		return (*(pfk_parse_inode_ftable[which_pfe]))
 				(bio, inode, key_info, algo_mode, is_pfe,
-					data_unit, (const char *)s_type);
+					(const char *)s_type);
 	}
 
 	/*
diff --git a/security/pfe/pfk_ext4.c b/security/pfe/pfk_ext4.c
index 7000b66..6d3cd4c 100644
--- a/security/pfe/pfk_ext4.c
+++ b/security/pfe/pfk_ext4.c
@@ -142,7 +142,6 @@
 	struct pfk_key_info *key_info,
 	enum ice_cryto_algo_mode *algo,
 	bool *is_pfe,
-	unsigned int *data_unit,
 	const char *storage_type)
 {
 	int ret = 0;
@@ -157,19 +156,6 @@
 	 */
 	*is_pfe = true;
 
-	/* Update dun based upon storage type.
-	 * For ext4 FS UFS has 4k dun whereas eMMC
-	 * uses 512Byte dun.
-	 */
-	if (storage_type && data_unit) {
-		if (!memcmp(storage_type, "ufs", strlen("ufs")))
-			*data_unit = 1 << ICE_CRYPTO_DATA_UNIT_4_KB;
-		else if (!memcmp(storage_type, "sdcc", strlen("sdcc")))
-			*data_unit = 1 << ICE_CRYPTO_DATA_UNIT_512_B;
-		else
-			return -EINVAL;
-	}
-
 	if (!pfk_ext4_is_ready())
 		return -ENODEV;
 
diff --git a/security/pfe/pfk_ext4.h b/security/pfe/pfk_ext4.h
index e39d04d..346027b 100644
--- a/security/pfe/pfk_ext4.h
+++ b/security/pfe/pfk_ext4.h
@@ -25,7 +25,6 @@
 	struct pfk_key_info *key_info,
 	enum ice_cryto_algo_mode *algo,
 	bool *is_pfe,
-	unsigned int *data_unit,
 	const char *storage_type);
 
 bool pfk_ext4_allow_merge_bio(const struct bio *bio1,
diff --git a/security/pfe/pfk_f2fs.c b/security/pfe/pfk_f2fs.c
index 2076267..9523c35 100644
--- a/security/pfe/pfk_f2fs.c
+++ b/security/pfe/pfk_f2fs.c
@@ -117,7 +117,6 @@
 		struct pfk_key_info *key_info,
 		enum ice_cryto_algo_mode *algo,
 		bool *is_pfe,
-		unsigned int *data_unit,
 		const char *storage_type)
 {
 	int ret = 0;
@@ -132,18 +131,6 @@
 	 */
 	*is_pfe = true;
 
-	/* Update the dun based upon storage type.
-	 * Right now both UFS and eMMC storage uses 4KB dun
-	 * for F2FS
-	 */
-	if (storage_type && data_unit) {
-		if (!memcmp(storage_type, "ufs", strlen("ufs")) ||
-			!memcmp(storage_type, "sdcc", strlen("sdcc")))
-			*data_unit = 1 << ICE_CRYPTO_DATA_UNIT_4_KB;
-		else
-			return -EINVAL;
-	}
-
 	if (!pfk_f2fs_is_ready())
 		return -ENODEV;
 
diff --git a/security/pfe/pfk_f2fs.h b/security/pfe/pfk_f2fs.h
index 2e0c21d..1b00bd8 100644
--- a/security/pfe/pfk_f2fs.h
+++ b/security/pfe/pfk_f2fs.h
@@ -25,7 +25,6 @@
 		struct pfk_key_info *key_info,
 		enum ice_cryto_algo_mode *algo,
 		bool *is_pfe,
-		unsigned int *data_unit,
 		const char *storage_type);
 
 bool pfk_f2fs_allow_merge_bio(const struct bio *bio1,