#include <linux/reset-controller.h>
 #include <linux/arm-smccc.h>
 
+#include <asm/smp_plat.h>
+
 #include "qcom_scm.h"
 
 static bool download_mode = IS_ENABLED(CONFIG_QCOM_SCM_DOWNLOAD_MODE_DEFAULT);
        return ret ? false : !!res.result[0];
 }
 
-/**
- * qcom_scm_set_warm_boot_addr() - Set the warm boot address for cpus
- * @entry: Entry point function for the cpus
- * @cpus: The cpumask of cpus that will use the entry point
- *
- * Set the Linux entry point for the SCM to transfer control to when coming
- * out of a power down. CPU power down may be executed on cpuidle or hotplug.
- */
-int qcom_scm_set_warm_boot_addr(void *entry, const cpumask_t *cpus)
+static int __qcom_scm_set_boot_addr_mc(void *entry, const cpumask_t *cpus,
+                                      unsigned int flags)
+{
+       struct qcom_scm_desc desc = {
+               .svc = QCOM_SCM_SVC_BOOT,
+               .cmd = QCOM_SCM_BOOT_SET_ADDR_MC,
+               .owner = ARM_SMCCC_OWNER_SIP,
+               .arginfo = QCOM_SCM_ARGS(6),
+       };
+       unsigned int cpu;
+       u64 map;
+
+       /* Need a device for DMA of the additional arguments */
+       if (!__scm || __get_convention() == SMC_CONVENTION_LEGACY)
+               return -EOPNOTSUPP;
+
+       desc.args[0] = virt_to_phys(entry);
+       for_each_cpu(cpu, cpus) {
+               map = cpu_logical_map(cpu);
+               desc.args[1] |= BIT(MPIDR_AFFINITY_LEVEL(map, 0));
+               desc.args[2] |= BIT(MPIDR_AFFINITY_LEVEL(map, 1));
+               desc.args[3] |= BIT(MPIDR_AFFINITY_LEVEL(map, 2));
+       }
+       desc.args[4] = ~0ULL; /* Reserved for affinity level 3 */
+       desc.args[5] = flags;
+
+       return qcom_scm_call(__scm->dev, &desc, NULL);
+}
+
+static int __qcom_scm_set_warm_boot_addr(void *entry, const cpumask_t *cpus)
 {
        int ret;
        int flags = 0;
 
        return ret;
 }
-EXPORT_SYMBOL(qcom_scm_set_warm_boot_addr);
 
 /**
- * qcom_scm_set_cold_boot_addr() - Set the cold boot address for cpus
+ * qcom_scm_set_warm_boot_addr() - Set the warm boot address for cpus
  * @entry: Entry point function for the cpus
  * @cpus: The cpumask of cpus that will use the entry point
  *
- * Set the cold boot address of the cpus. Any cpu outside the supported
- * range would be removed from the cpu present mask.
+ * Set the Linux entry point for the SCM to transfer control to when coming
+ * out of a power down. CPU power down may be executed on cpuidle or hotplug.
  */
-int qcom_scm_set_cold_boot_addr(void *entry, const cpumask_t *cpus)
+int qcom_scm_set_warm_boot_addr(void *entry, const cpumask_t *cpus)
+{
+       if (!cpus || cpumask_empty(cpus))
+               return -EINVAL;
+
+       if (__qcom_scm_set_boot_addr_mc(entry, cpus, QCOM_SCM_BOOT_MC_FLAG_WARMBOOT))
+               /* Fallback to old SCM call */
+               return __qcom_scm_set_warm_boot_addr(entry, cpus);
+       return 0;
+}
+EXPORT_SYMBOL(qcom_scm_set_warm_boot_addr);
+
+static int __qcom_scm_set_cold_boot_addr(void *entry, const cpumask_t *cpus)
 {
        int flags = 0;
        int cpu;
                .owner = ARM_SMCCC_OWNER_SIP,
        };
 
-       if (!cpus || cpumask_empty(cpus))
-               return -EINVAL;
-
        for_each_cpu(cpu, cpus) {
                if (cpu < ARRAY_SIZE(scm_cb_flags))
                        flags |= scm_cb_flags[cpu];
 
        return qcom_scm_call_atomic(__scm ? __scm->dev : NULL, &desc, NULL);
 }
+
+/**
+ * qcom_scm_set_cold_boot_addr() - Set the cold boot address for cpus
+ * @entry: Entry point function for the cpus
+ * @cpus: The cpumask of cpus that will use the entry point
+ *
+ * Set the cold boot address of the cpus. Any cpu outside the supported
+ * range would be removed from the cpu present mask.
+ */
+int qcom_scm_set_cold_boot_addr(void *entry, const cpumask_t *cpus)
+{
+       if (!cpus || cpumask_empty(cpus))
+               return -EINVAL;
+
+       if (__qcom_scm_set_boot_addr_mc(entry, cpus, QCOM_SCM_BOOT_MC_FLAG_COLDBOOT))
+               /* Fallback to old SCM call */
+               return __qcom_scm_set_cold_boot_addr(entry, cpus);
+       return 0;
+}
 EXPORT_SYMBOL(qcom_scm_set_cold_boot_addr);
 
 /**