diff --git a/drivers/scsi/ufs/ufs-tegra.c b/drivers/scsi/ufs/ufs-tegra.c
index 780ee83d..3ae829a8 100644
--- a/drivers/scsi/ufs/ufs-tegra.c
+++ b/drivers/scsi/ufs/ufs-tegra.c
@@ -17,6 +17,7 @@
 #include <linux/of_device.h>
 #include <linux/of_platform.h>
 #include <linux/of_gpio.h>
+#include <linux/iommu.h>
 
 #ifdef CONFIG_DEBUG_FS
 #include <linux/debugfs.h>
@@ -1063,6 +1064,13 @@ static int ufs_tegra_resume(struct ufs_hba *hba, enum ufs_pm_op pm_op)
 		 * T234 does not require context restore
 		 */
 		ufs_tegra_context_restore(ufs_tegra);
+	} else {
+		writel(UFS_AUX_ADDR_VIRT_CTRL_EN,
+			ufs_tegra->ufs_virtualization_base +
+			UFS_AUX_ADDR_VIRT_CTRL_0);
+		writel(ufs_tegra->streamid,
+			ufs_tegra->ufs_virtualization_base +
+			UFS_AUX_ADDR_VIRT_REG_0);
 	}
 
 	ufs_tegra_set_clk_div(hba, UFS_VNDR_HCLKDIV_1US_TICK);
@@ -1390,6 +1398,9 @@ static void ufs_tegra_config_soc_data(struct ufs_tegra_host *ufs_tegra)
 
 	ufs_tegra->enable_scramble =
 		of_property_read_bool(np, "nvidia,enable-scramble");
+
+	if (ufs_tegra->soc->chip_id == TEGRA234)
+		ufs_tegra->hba->quirks |= UFSHCD_QUIRK_ENABLE_STREAM_ID;
 }
 
 static void ufs_tegra_eq_timeout(struct ufs_tegra_host *ufs_tegra)
@@ -1421,6 +1432,7 @@ static int ufs_tegra_init(struct ufs_hba *hba)
 	struct device *dev = hba->dev;
 	int err = 0;
 	resource_size_t ufs_aux_base_addr, ufs_aux_addr_range, mphy_addr_range;
+	struct iommu_fwspec *fwspec;
 
 	ufs_tegra = devm_kzalloc(dev, sizeof(*ufs_tegra), GFP_KERNEL);
 	if (!ufs_tegra) {
@@ -1470,6 +1482,17 @@ static int ufs_tegra_init(struct ufs_hba *hba)
 		goto out;
 	}
 
+	if (ufs_tegra->soc->chip_id == TEGRA234) {
+		ufs_tegra->ufs_virtualization_base = devm_ioremap(dev,
+				NV_ADDRESS_MAP_T23X_UFSHC_VIRT_BASE,
+				UFS_AUX_ADDR_VIRT_RANGE_23X);
+		if (!ufs_tegra->ufs_virtualization_base) {
+			err = -ENOMEM;
+			dev_err(dev, "UFS Virtualization failed\n");
+			goto out;
+		}
+	}
+
 	err = ufs_tegra_init_ufs_clks(ufs_tegra);
 	if (err)
 		goto out_host_free;
@@ -1492,6 +1515,23 @@ static int ufs_tegra_init(struct ufs_hba *hba)
 	if (err)
 		goto out_host_free;
 
+	if (ufs_tegra->soc->chip_id == TEGRA234) {
+		fwspec = dev_iommu_fwspec_get(dev);
+		if (fwspec == NULL) {
+			err = -ENODEV;
+			dev_err(dev, "Failed to get MC streamidd\n");
+			goto out;
+		} else {
+			ufs_tegra->streamid = fwspec->ids[0] & 0xffff;
+			writel(UFS_AUX_ADDR_VIRT_CTRL_EN,
+				ufs_tegra->ufs_virtualization_base +
+				UFS_AUX_ADDR_VIRT_CTRL_0);
+			writel(ufs_tegra->streamid,
+				ufs_tegra->ufs_virtualization_base +
+				UFS_AUX_ADDR_VIRT_REG_0);
+		}
+	}
+
 	err = ufs_tegra_enable_mphylane_clks(ufs_tegra);
 	if (err)
 		goto out_disable_ufs_clks;
diff --git a/drivers/scsi/ufs/ufs-tegra.h b/drivers/scsi/ufs/ufs-tegra.h
index c1e0516c..4e447670 100644
--- a/drivers/scsi/ufs/ufs-tegra.h
+++ b/drivers/scsi/ufs/ufs-tegra.h
@@ -9,8 +9,10 @@
 
 #define NV_ADDRESS_MAP_MPHY_L0_BASE		0x02470000
 #define NV_ADDRESS_MAP_MPHY_L1_BASE		0x02480000
+#define NV_ADDRESS_MAP_T23X_UFSHC_VIRT_BASE	0x02520000
 #define MPHY_ADDR_RANGE_T234			0x2268
 #define MPHY_ADDR_RANGE				0x200
+#define UFS_AUX_ADDR_VIRT_RANGE_23X		0x14f
 
 /* UFS AUX Base address for T194 */
 #define NV_ADDRESS_MAP_UFSHC_AUX_BASE		0x02460000
@@ -21,7 +23,9 @@
 #define NV_ADDRESS_MAP_T23X_UFSHC_AUX_BASE	0x02510000
 /* UFS AUX address range in T234 */
 #define UFS_AUX_ADDR_RANGE_23X			0x20
-
+#define UFS_AUX_ADDR_VIRT_CTRL_0		0x0
+#define UFS_AUX_ADDR_VIRT_CTRL_EN		0x1
+#define UFS_AUX_ADDR_VIRT_REG_0			0x4
 
 /*
  * M-PHY Registers
@@ -320,6 +324,7 @@ struct ufs_tegra_host {
 	void __iomem *mphy_l0_base;
 	void __iomem *mphy_l1_base;
 	void __iomem *ufs_aux_base;
+	void __iomem *ufs_virtualization_base;
 	struct reset_control *ufs_rst;
 	struct reset_control *ufs_axi_m_rst;
 	struct reset_control *ufshc_lp_rst;
@@ -366,6 +371,7 @@ struct ufs_tegra_host {
 	bool enable_scramble;
 	u32 ref_clk_freq;
 	struct ufs_tegra_soc_data *soc;
+	u32 streamid;
 #ifdef CONFIG_DEBUG_FS
 	u32 refclk_value;
 	long program_refclk;
diff --git a/include/drivers-private/scsi/ufs/k515/ufshcd.h b/include/drivers-private/scsi/ufs/k515/ufshcd.h
index 07ada667..1167b7a0 100644
--- a/include/drivers-private/scsi/ufs/k515/ufshcd.h
+++ b/include/drivers-private/scsi/ufs/k515/ufshcd.h
@@ -588,6 +588,11 @@ enum ufshcd_quirks {
 	 * This quirk allows only sg entries aligned with page size.
 	 */
 	UFSHCD_QUIRK_ALIGN_SG_WITH_PAGE_SIZE		= 1 << 14,
+
+	/*
+	 * Enable this quirk to support Stream-ID programming
+	 */
+	UFSHCD_QUIRK_ENABLE_STREAM_ID			= 1 << 15,
 };
 
 enum ufshcd_caps {
diff --git a/include/drivers-private/scsi/ufs/k516/ufshcd.h b/include/drivers-private/scsi/ufs/k516/ufshcd.h
index 88c20f36..21d33614 100644
--- a/include/drivers-private/scsi/ufs/k516/ufshcd.h
+++ b/include/drivers-private/scsi/ufs/k516/ufshcd.h
@@ -602,6 +602,11 @@ enum ufshcd_quirks {
 	 * support physical host configuration.
 	 */
 	UFSHCD_QUIRK_SKIP_PH_CONFIGURATION		= 1 << 16,
+
+	/*
+	 * Enable this quirk to support Stream-ID programming
+	 */
+	UFSHCD_QUIRK_ENABLE_STREAM_ID			= 1 << 17,
 };
 
 enum ufshcd_caps {