{
        int ret;
        u64 smc_arg[2];
+       u32 module_id;
+       u32 feature_check_api_id;
 
-       smc_arg[0] = PM_SIP_SVC | PM_FEATURE_CHECK;
-       smc_arg[1] = api_id;
+       module_id = FIELD_GET(MODULE_ID_MASK, api_id);
+
+       /*
+        * Feature check of APIs belonging to PM, XSEM, and TF-A are handled by calling
+        * PM_FEATURE_CHECK API. For other modules, call PM_API_FEATURES API.
+        */
+       if (module_id == PM_MODULE_ID || module_id == XSEM_MODULE_ID || module_id == TF_A_MODULE_ID)
+               feature_check_api_id = PM_FEATURE_CHECK;
+       else
+               feature_check_api_id = PM_API_FEATURES;
+
+       /*
+        * Feature check of TF-A APIs is done in the TF-A layer and it expects for
+        * MODULE_ID_MASK bits of SMC's arg[0] to be the same as PM_MODULE_ID.
+        */
+       if (module_id == TF_A_MODULE_ID)
+               module_id = PM_MODULE_ID;
+
+       smc_arg[0] = PM_SIP_SVC | FIELD_PREP(MODULE_ID_MASK, module_id) | feature_check_api_id;
+       smc_arg[1] = (api_id & API_ID_MASK);
 
        ret = do_fw_call(ret_payload, 2, smc_arg[0], smc_arg[1]);
        if (ret)
        if (ret)
                return ret;
 
-       np = of_find_compatible_node(NULL, NULL, "xlnx,zynqmp");
-       if (!np) {
-               np = of_find_compatible_node(NULL, NULL, "xlnx,versal");
-               if (!np)
-                       return 0;
-
+       ret = do_feature_check_call(PM_FEATURE_CHECK);
+       if (ret >= 0 && ((ret & FIRMWARE_VERSION_MASK) >= PM_API_VERSION_1))
                feature_check_enabled = true;
-       }
-
-       if (!feature_check_enabled) {
-               ret = do_feature_check_call(PM_FEATURE_CHECK);
-               if (ret >= 0)
-                       feature_check_enabled = true;
-       }
-
-       of_node_put(np);
 
        devinfo = devm_kzalloc(dev, sizeof(*devinfo), GFP_KERNEL);
        if (!devinfo)
 
 #define PM_SIP_SVC                     0xC2000000
 
 /* PM API versions */
+#define PM_API_VERSION_1       1
 #define PM_API_VERSION_2       2
 
 #define PM_PINCTRL_PARAM_SET_VERSION   2
 #define FAMILY_CODE_MASK       GENMASK(27, 21)
 #define SUB_FAMILY_CODE_MASK   GENMASK(20, 19)
 
+#define API_ID_MASK            GENMASK(7, 0)
+#define MODULE_ID_MASK         GENMASK(11, 8)
+
 /* ATF only commands */
 #define TF_A_PM_REGISTER_SGI           0xa04
 #define PM_GET_TRUSTZONE_VERSION       0xa03
 #define XPM_EVENT_ERROR_MASK_NOC_NCR           BIT(13)
 #define XPM_EVENT_ERROR_MASK_NOC_CR            BIT(12)
 
+enum pm_module_id {
+       PM_MODULE_ID = 0x0,
+       XSEM_MODULE_ID = 0x3,
+       TF_A_MODULE_ID = 0xa,
+};
+
 enum pm_api_cb_id {
        PM_INIT_SUSPEND_CB = 30,
        PM_ACKNOWLEDGE_CB = 31,
 };
 
 enum pm_api_id {
+       PM_API_FEATURES = 0,
        PM_GET_API_VERSION = 1,
        PM_REGISTER_NOTIFIER = 5,
        PM_FORCE_POWERDOWN = 8,