diff --git a/drivers/video/omap2/dss/core.c b/drivers/video/omap2/dss/core.c
index b2af72dc20bdb2cf0cc5f880008916fb85d09826..d94ef9e31a354b877f2eb3dcee6cb4be1bda2462 100644
--- a/drivers/video/omap2/dss/core.c
+++ b/drivers/video/omap2/dss/core.c
@@ -237,7 +237,7 @@ static int __init omap_dss_probe(struct platform_device *pdev)
 
 	core.pdev = pdev;
 
-	dss_features_init();
+	dss_features_init(pdata->version);
 
 	dss_apply_init();
 
diff --git a/drivers/video/omap2/dss/dss_features.c b/drivers/video/omap2/dss/dss_features.c
index acbc1e1efba3650aaccba2d86b000816973d23ba..5936ba78c14adbc1b90b7656995dcd0c25cef2a3 100644
--- a/drivers/video/omap2/dss/dss_features.c
+++ b/drivers/video/omap2/dss/dss_features.c
@@ -929,29 +929,44 @@ bool dss_feat_rotation_type_supported(enum omap_dss_rotation_type rot_type)
 	return omap_current_dss_features->supported_rotation_types & rot_type;
 }
 
-void dss_features_init(void)
+void dss_features_init(enum omapdss_version version)
 {
-	if (cpu_is_omap24xx())
+	switch (version) {
+	case OMAPDSS_VER_OMAP24xx:
 		omap_current_dss_features = &omap2_dss_features;
-	else if (cpu_is_omap3630())
+		break;
+
+	case OMAPDSS_VER_OMAP34xx_ES1:
+	case OMAPDSS_VER_OMAP34xx_ES3:
+		omap_current_dss_features = &omap3430_dss_features;
+		break;
+
+	case OMAPDSS_VER_OMAP3630:
 		omap_current_dss_features = &omap3630_dss_features;
-	else if (cpu_is_omap34xx()) {
-		if (soc_is_am35xx()) {
-			omap_current_dss_features = &am35xx_dss_features;
-		} else {
-			omap_current_dss_features = &omap3430_dss_features;
-		}
-	}
-	else if (omap_rev() == OMAP4430_REV_ES1_0)
+		break;
+
+	case OMAPDSS_VER_OMAP4430_ES1:
 		omap_current_dss_features = &omap4430_es1_0_dss_features;
-	else if (omap_rev() == OMAP4430_REV_ES2_0 ||
-		omap_rev() == OMAP4430_REV_ES2_1 ||
-		omap_rev() == OMAP4430_REV_ES2_2)
+		break;
+
+	case OMAPDSS_VER_OMAP4430_ES2:
 		omap_current_dss_features = &omap4430_es2_0_1_2_dss_features;
-	else if (cpu_is_omap44xx())
+		break;
+
+	case OMAPDSS_VER_OMAP4:
 		omap_current_dss_features = &omap4_dss_features;
-	else if (soc_is_omap54xx())
+		break;
+
+	case OMAPDSS_VER_OMAP5:
 		omap_current_dss_features = &omap5_dss_features;
-	else
+		break;
+
+	case OMAPDSS_VER_AM35xx:
+		omap_current_dss_features = &am35xx_dss_features;
+		break;
+
+	default:
 		DSSWARN("Unsupported OMAP version");
+		break;
+	}
 }
diff --git a/drivers/video/omap2/dss/dss_features.h b/drivers/video/omap2/dss/dss_features.h
index 9218113b5e88c2c5488acb2f989e06ff080cfe06..14a412ea8dc6cae56dd93a9c10dac78eb340dcf4 100644
--- a/drivers/video/omap2/dss/dss_features.h
+++ b/drivers/video/omap2/dss/dss_features.h
@@ -123,7 +123,7 @@ bool dss_feat_rotation_type_supported(enum omap_dss_rotation_type rot_type);
 
 bool dss_has_feature(enum dss_feat_id id);
 void dss_feat_get_reg_field(enum dss_feat_reg_field id, u8 *start, u8 *end);
-void dss_features_init(void);
+void dss_features_init(enum omapdss_version version);
 #if defined(CONFIG_OMAP4_DSS_HDMI)
 void dss_init_hdmi_ip_ops(struct hdmi_ip_data *ip_data);
 #endif