[5/7] drivers: firmware: psci: Simplify state node parsing

Message ID 20190228135919.3747-6-ulf.hansson@linaro.org
State New
Headers show
Series
  • drivers: firmware: psci: Some cleanup and refactoring
Related show

Commit Message

Ulf Hansson Feb. 28, 2019, 1:59 p.m.
Instead of iterating through all the state nodes in DT, to find out how
many states that needs to be allocated, let's use the number already known
by the cpuidle driver. In this way we can drop the iteration altogether.

Signed-off-by: Ulf Hansson <ulf.hansson@linaro.org>

---
 drivers/firmware/psci/psci.c | 25 ++++++++++++-------------
 1 file changed, 12 insertions(+), 13 deletions(-)

-- 
2.17.1

Comments

Daniel Lezcano Feb. 28, 2019, 3:40 p.m. | #1
On 28/02/2019 14:59, Ulf Hansson wrote:
> Instead of iterating through all the state nodes in DT, to find out how

> many states that needs to be allocated, let's use the number already known

> by the cpuidle driver. In this way we can drop the iteration altogether.

> 

> Signed-off-by: Ulf Hansson <ulf.hansson@linaro.org>

> ---

>  drivers/firmware/psci/psci.c | 25 ++++++++++++-------------

>  1 file changed, 12 insertions(+), 13 deletions(-)

> 

> diff --git a/drivers/firmware/psci/psci.c b/drivers/firmware/psci/psci.c

> index d50b46a0528f..cbfc936d251c 100644

> --- a/drivers/firmware/psci/psci.c

> +++ b/drivers/firmware/psci/psci.c

> @@ -290,26 +290,20 @@ static int psci_dt_parse_state_node(struct device_node *np, u32 *state)

>  static int psci_dt_cpu_init_idle(struct cpuidle_driver *drv,

>  			struct device_node *cpu_node, int cpu)

>  {

> -	int i, ret = 0, count = 0;

> +	int i, ret = 0, num_state_nodes = drv->state_count - 1;


why -1 ?

>  	u32 *psci_states;

>  	struct device_node *state_node;

>  

> -	/* Count idle states */

> -	while ((state_node = of_parse_phandle(cpu_node, "cpu-idle-states",

> -					      count))) {

> -		count++;

> -		of_node_put(state_node);

> -	}

> -

> -	if (!count)

> -		return -ENODEV;

> -

> -	psci_states = kcalloc(count, sizeof(*psci_states), GFP_KERNEL);

> +	psci_states = kcalloc(num_state_nodes, sizeof(*psci_states),

> +			GFP_KERNEL);

>  	if (!psci_states)

>  		return -ENOMEM;

>  

> -	for (i = 0; i < count; i++) {

> +	for (i = 0; i < num_state_nodes; i++) {

>  		state_node = of_parse_phandle(cpu_node, "cpu-idle-states", i);

> +		if (!state_node)

> +			break;

> +

>  		ret = psci_dt_parse_state_node(state_node, &psci_states[i]);

>  		of_node_put(state_node);

>  

> @@ -319,6 +313,11 @@ static int psci_dt_cpu_init_idle(struct cpuidle_driver *drv,

>  		pr_debug("psci-power-state %#x index %d\n", psci_states[i], i);

>  	}

>  

> +	if (i != num_state_nodes) {

> +		ret = -ENODEV;

> +		goto free_mem;

> +	}

> +

>  	/* Idle states parsed correctly, initialize per-cpu pointer */

>  	per_cpu(psci_power_state, cpu) = psci_states;

>  	return 0;

> 



-- 
 <http://www.linaro.org/> Linaro.org │ Open source software for ARM SoCs

Follow Linaro:  <http://www.facebook.com/pages/Linaro> Facebook |
<http://twitter.com/#!/linaroorg> Twitter |
<http://www.linaro.org/linaro-blog/> Blog
Mark Rutland March 1, 2019, 5:28 p.m. | #2
On Thu, Feb 28, 2019 at 02:59:17PM +0100, Ulf Hansson wrote:
> Instead of iterating through all the state nodes in DT, to find out how

> many states that needs to be allocated, let's use the number already known

> by the cpuidle driver. In this way we can drop the iteration altogether.

> 

> Signed-off-by: Ulf Hansson <ulf.hansson@linaro.org>

> ---

>  drivers/firmware/psci/psci.c | 25 ++++++++++++-------------

>  1 file changed, 12 insertions(+), 13 deletions(-)

> 

> diff --git a/drivers/firmware/psci/psci.c b/drivers/firmware/psci/psci.c

> index d50b46a0528f..cbfc936d251c 100644

> --- a/drivers/firmware/psci/psci.c

> +++ b/drivers/firmware/psci/psci.c

> @@ -290,26 +290,20 @@ static int psci_dt_parse_state_node(struct device_node *np, u32 *state)

>  static int psci_dt_cpu_init_idle(struct cpuidle_driver *drv,

>  			struct device_node *cpu_node, int cpu)

>  {

> -	int i, ret = 0, count = 0;

> +	int i, ret = 0, num_state_nodes = drv->state_count - 1;

>  	u32 *psci_states;

>  	struct device_node *state_node;

>  

> -	/* Count idle states */

> -	while ((state_node = of_parse_phandle(cpu_node, "cpu-idle-states",

> -					      count))) {

> -		count++;

> -		of_node_put(state_node);

> -	}

> -


To be honest, I'd rather not tighten the coupling with the cpuidle
driver here. For example, I'm not that happy with the PSCI backend
having to know the driver has a specific WFI state.

IIUC we could get rid of the explicit counting with something like:

	count = of_parse_phandle_with_args(cpu_node, "cpu-idle-states", NULL);

... but I'm not sure that the overall change is a simplification.

Does this change make it easier to plumb in something in future?

Thanks,
Mark.

> -	if (!count)

> -		return -ENODEV;

> -

> -	psci_states = kcalloc(count, sizeof(*psci_states), GFP_KERNEL);

> +	psci_states = kcalloc(num_state_nodes, sizeof(*psci_states),

> +			GFP_KERNEL);

>  	if (!psci_states)

>  		return -ENOMEM;

>  

> -	for (i = 0; i < count; i++) {

> +	for (i = 0; i < num_state_nodes; i++) {

>  		state_node = of_parse_phandle(cpu_node, "cpu-idle-states", i);

> +		if (!state_node)

> +			break;

> +

>  		ret = psci_dt_parse_state_node(state_node, &psci_states[i]);

>  		of_node_put(state_node);

>  

> @@ -319,6 +313,11 @@ static int psci_dt_cpu_init_idle(struct cpuidle_driver *drv,

>  		pr_debug("psci-power-state %#x index %d\n", psci_states[i], i);

>  	}

>  

> +	if (i != num_state_nodes) {

> +		ret = -ENODEV;

> +		goto free_mem;

> +	}

> +

>  	/* Idle states parsed correctly, initialize per-cpu pointer */

>  	per_cpu(psci_power_state, cpu) = psci_states;

>  	return 0;

> -- 

> 2.17.1

>
Ulf Hansson March 4, 2019, 10:14 a.m. | #3
On Fri, 1 Mar 2019 at 18:28, Mark Rutland <mark.rutland@arm.com> wrote:
>

> On Thu, Feb 28, 2019 at 02:59:17PM +0100, Ulf Hansson wrote:

> > Instead of iterating through all the state nodes in DT, to find out how

> > many states that needs to be allocated, let's use the number already known

> > by the cpuidle driver. In this way we can drop the iteration altogether.

> >

> > Signed-off-by: Ulf Hansson <ulf.hansson@linaro.org>

> > ---

> >  drivers/firmware/psci/psci.c | 25 ++++++++++++-------------

> >  1 file changed, 12 insertions(+), 13 deletions(-)

> >

> > diff --git a/drivers/firmware/psci/psci.c b/drivers/firmware/psci/psci.c

> > index d50b46a0528f..cbfc936d251c 100644

> > --- a/drivers/firmware/psci/psci.c

> > +++ b/drivers/firmware/psci/psci.c

> > @@ -290,26 +290,20 @@ static int psci_dt_parse_state_node(struct device_node *np, u32 *state)

> >  static int psci_dt_cpu_init_idle(struct cpuidle_driver *drv,

> >                       struct device_node *cpu_node, int cpu)

> >  {

> > -     int i, ret = 0, count = 0;

> > +     int i, ret = 0, num_state_nodes = drv->state_count - 1;

> >       u32 *psci_states;

> >       struct device_node *state_node;

> >

> > -     /* Count idle states */

> > -     while ((state_node = of_parse_phandle(cpu_node, "cpu-idle-states",

> > -                                           count))) {

> > -             count++;

> > -             of_node_put(state_node);

> > -     }

> > -

>

> To be honest, I'd rather not tighten the coupling with the cpuidle

> driver here. For example, I'm not that happy with the PSCI backend

> having to know the driver has a specific WFI state.


If you ask me, the coupling is already there, only that it's hidden by
taking assumptions about the WFI state in the code.

However, I certainly agree with you, that this isn't very nice.

>

> IIUC we could get rid of the explicit counting with something like:

>

>         count = of_parse_phandle_with_args(cpu_node, "cpu-idle-states", NULL);

>

> ... but I'm not sure that the overall change is a simplification.


In my opinion, no it doesn't.

To me, it seems a kind of silly (and in-efficient) to do an OF parsing
that has already been done and given the information we need.

>

> Does this change make it easier to plumb in something in future?


Yes, I need this for additional changes on top.

Note that, patch4 also provides the opportunity to do a similar
cleanup of the initialization code in drivers/soc/qcom/spm.c. I
haven't made that part of this series though.

I guess in the end, we need to accept that part of the psci driver is
really a cpuidle driver. Trying to keep them entirely separate,
doesn't come without complexity/churns.

While working on psci changes in the recent series I have posted, I
was even considering adding a completely new cpuidle driver for psci
(in drivers/cpuidle/*) and instead define a number of psci interface
functions, which that driver could call. That would probably be a
better split, but requires quite some changes.

So, what do you want me to do with this?

>

> Thanks,

> Mark.


Thanks a lot for reviewing!

Kind regards
Uffe

>

> > -     if (!count)

> > -             return -ENODEV;

> > -

> > -     psci_states = kcalloc(count, sizeof(*psci_states), GFP_KERNEL);

> > +     psci_states = kcalloc(num_state_nodes, sizeof(*psci_states),

> > +                     GFP_KERNEL);

> >       if (!psci_states)

> >               return -ENOMEM;

> >

> > -     for (i = 0; i < count; i++) {

> > +     for (i = 0; i < num_state_nodes; i++) {

> >               state_node = of_parse_phandle(cpu_node, "cpu-idle-states", i);

> > +             if (!state_node)

> > +                     break;

> > +

> >               ret = psci_dt_parse_state_node(state_node, &psci_states[i]);

> >               of_node_put(state_node);

> >

> > @@ -319,6 +313,11 @@ static int psci_dt_cpu_init_idle(struct cpuidle_driver *drv,

> >               pr_debug("psci-power-state %#x index %d\n", psci_states[i], i);

> >       }

> >

> > +     if (i != num_state_nodes) {

> > +             ret = -ENODEV;

> > +             goto free_mem;

> > +     }

> > +

> >       /* Idle states parsed correctly, initialize per-cpu pointer */

> >       per_cpu(psci_power_state, cpu) = psci_states;

> >       return 0;

> > --

> > 2.17.1

> >
Lorenzo Pieralisi March 6, 2019, 6:15 p.m. | #4
On Mon, Mar 04, 2019 at 11:14:18AM +0100, Ulf Hansson wrote:
> On Fri, 1 Mar 2019 at 18:28, Mark Rutland <mark.rutland@arm.com> wrote:

> >

> > On Thu, Feb 28, 2019 at 02:59:17PM +0100, Ulf Hansson wrote:

> > > Instead of iterating through all the state nodes in DT, to find out how

> > > many states that needs to be allocated, let's use the number already known

> > > by the cpuidle driver. In this way we can drop the iteration altogether.

> > >

> > > Signed-off-by: Ulf Hansson <ulf.hansson@linaro.org>

> > > ---

> > >  drivers/firmware/psci/psci.c | 25 ++++++++++++-------------

> > >  1 file changed, 12 insertions(+), 13 deletions(-)

> > >

> > > diff --git a/drivers/firmware/psci/psci.c b/drivers/firmware/psci/psci.c

> > > index d50b46a0528f..cbfc936d251c 100644

> > > --- a/drivers/firmware/psci/psci.c

> > > +++ b/drivers/firmware/psci/psci.c

> > > @@ -290,26 +290,20 @@ static int psci_dt_parse_state_node(struct device_node *np, u32 *state)

> > >  static int psci_dt_cpu_init_idle(struct cpuidle_driver *drv,

> > >                       struct device_node *cpu_node, int cpu)

> > >  {

> > > -     int i, ret = 0, count = 0;

> > > +     int i, ret = 0, num_state_nodes = drv->state_count - 1;

> > >       u32 *psci_states;

> > >       struct device_node *state_node;

> > >

> > > -     /* Count idle states */

> > > -     while ((state_node = of_parse_phandle(cpu_node, "cpu-idle-states",

> > > -                                           count))) {

> > > -             count++;

> > > -             of_node_put(state_node);

> > > -     }

> > > -

> >

> > To be honest, I'd rather not tighten the coupling with the cpuidle

> > driver here. For example, I'm not that happy with the PSCI backend

> > having to know the driver has a specific WFI state.

> 

> If you ask me, the coupling is already there, only that it's hidden by

> taking assumptions about the WFI state in the code.


There is no assumption taken - I wrote it down here:

https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/tree/Documentation/devicetree/bindings/arm/idle-states.txt?h=v5.0


> However, I certainly agree with you, that this isn't very nice.


The idea behind the generic ARM CPU idle driver is as follows:

- A generic front-end in drivers/cpuidle/cpuidle-arm.c
- An arch back-end (that is defined by the enable-method), on ARM64
  it is PSCI

As usual with the ARM CPUidle mess, there must be logic connecting
the front-end and the back-end. An idle state index was used since
I saw no other generic way. If there are better ideas they are welcome.

Otherwise we must go back to having a PSCI specific CPUIdle driver
and, on arch/arm, platform specific CPUidle drivers.

The aim was to simplify but to do that we need a connection logic
between drivers<->arch code, that's the only way we can have a generic
idle driver and corresponding boilerplate.

> > IIUC we could get rid of the explicit counting with something like:

> >

> >         count = of_parse_phandle_with_args(cpu_node, "cpu-idle-states", NULL);

> >

> > ... but I'm not sure that the overall change is a simplification.

> 

> In my opinion, no it doesn't.

> 

> To me, it seems a kind of silly (and in-efficient) to do an OF parsing

> that has already been done and given the information we need.


Yeah. It is boot path with idle states in the order of (max) dozens,
silly and inefficient, maybe but that should be fine.

See above.

> > Does this change make it easier to plumb in something in future?

> 

> Yes, I need this for additional changes on top.

> 

> Note that, patch4 also provides the opportunity to do a similar

> cleanup of the initialization code in drivers/soc/qcom/spm.c. I

> haven't made that part of this series though.

> 

> I guess in the end, we need to accept that part of the psci driver is

> really a cpuidle driver. Trying to keep them entirely separate,

> doesn't come without complexity/churns.


PSCI driver is a kernel interface to firmware, it is not a CPUidle
driver per-se, we tried to decouple firmware interfaces from kernel
data structures as much as possible, again, see above.

> While working on psci changes in the recent series I have posted, I

> was even considering adding a completely new cpuidle driver for psci

> (in drivers/cpuidle/*) and instead define a number of psci interface

> functions, which that driver could call. That would probably be a

> better split, but requires quite some changes.


It can be done but with it the whole generic ARM CPUidle driver
infrastructure must go with it (and you will still have a standard wfi
state in the PSCI idle state array anyway).

The idea behind ARM64 cpu_ops clashes a bit with this approach so
there is a discussion to be had here.

> So, what do you want me to do with this?


We need to answer the question above.

Thanks,
Lorenzo

> > Thanks,

> > Mark.

> 

> Thanks a lot for reviewing!

> 

> Kind regards

> Uffe

> 

> >

> > > -     if (!count)

> > > -             return -ENODEV;

> > > -

> > > -     psci_states = kcalloc(count, sizeof(*psci_states), GFP_KERNEL);

> > > +     psci_states = kcalloc(num_state_nodes, sizeof(*psci_states),

> > > +                     GFP_KERNEL);

> > >       if (!psci_states)

> > >               return -ENOMEM;

> > >

> > > -     for (i = 0; i < count; i++) {

> > > +     for (i = 0; i < num_state_nodes; i++) {

> > >               state_node = of_parse_phandle(cpu_node, "cpu-idle-states", i);

> > > +             if (!state_node)

> > > +                     break;

> > > +

> > >               ret = psci_dt_parse_state_node(state_node, &psci_states[i]);

> > >               of_node_put(state_node);

> > >

> > > @@ -319,6 +313,11 @@ static int psci_dt_cpu_init_idle(struct cpuidle_driver *drv,

> > >               pr_debug("psci-power-state %#x index %d\n", psci_states[i], i);

> > >       }

> > >

> > > +     if (i != num_state_nodes) {

> > > +             ret = -ENODEV;

> > > +             goto free_mem;

> > > +     }

> > > +

> > >       /* Idle states parsed correctly, initialize per-cpu pointer */

> > >       per_cpu(psci_power_state, cpu) = psci_states;

> > >       return 0;

> > > --

> > > 2.17.1

> > >
Lorenzo Pieralisi March 8, 2019, 1:17 p.m. | #5
On Fri, Mar 08, 2019 at 02:07:51PM +0100, Ulf Hansson wrote:
> On Fri, 8 Mar 2019 at 12:49, Lorenzo Pieralisi

> <lorenzo.pieralisi@arm.com> wrote:

> >

> > On Fri, Mar 08, 2019 at 11:36:49AM +0100, Ulf Hansson wrote:

> >

> > [...]

> >

> > > Instead, my suggestion is according to what I propose in patch 4 and

> > > $subject patch, which means minor adjustments to be able to pass the

> > > struct cpuidle_driver * to the init functions. This, I need it for

> > > next steps, but already at this point it improves things as it avoids

> > > some of the OF parsing, and that's good, isn't it?

> >

> > I will take the patches Mark ACKed and send them for v5.2 as

> > early as it gets in v5.1-rc* cycle.

> 

> Actually, may I suggest we funnel these through Rafael's tree, unless

> you are expecting other PSCI changes for v.5.2, which could cause

> conflicts?

> 

> The reason is, other PM core changes, to genpd for example, needs to

> go via Rafael's tree. Those would then potentially block us for

> applying any other changes to your tree (arm-soc?) for PSCI (as there

> is dependency) until v5.3.

> 

> How about if you provides your explicit acks for those PSCI changes

> your are happy with, then Rafael can pick them?


It is fine we can do that, I would have not sent the patches Mark
has ACKed to arm-soc till -{rc2/rc3} anyway.

Thanks,
Lorenzo

> > For this one maybe you can post the changes on top and see what's

> > the best way forward ?

> >

> > I agree that duplicating idle state parsing code across back-ends

> > is silly - we just want to keep PSCI and kernel data structure

> > decoupled.

> 

> Right. Let's continue this discussion later on and move forward here.

> Make sense.

> 

> >

> > Post the code on top and we will find a way forward, OK ?

> 

> Sure, let me do that.

> 

> Thanks and kind regards

> Uffe
Ulf Hansson March 8, 2019, 1:23 p.m. | #6
On Fri, 8 Mar 2019 at 14:17, Lorenzo Pieralisi
<lorenzo.pieralisi@arm.com> wrote:
>

> On Fri, Mar 08, 2019 at 02:07:51PM +0100, Ulf Hansson wrote:

> > On Fri, 8 Mar 2019 at 12:49, Lorenzo Pieralisi

> > <lorenzo.pieralisi@arm.com> wrote:

> > >

> > > On Fri, Mar 08, 2019 at 11:36:49AM +0100, Ulf Hansson wrote:

> > >

> > > [...]

> > >

> > > > Instead, my suggestion is according to what I propose in patch 4 and

> > > > $subject patch, which means minor adjustments to be able to pass the

> > > > struct cpuidle_driver * to the init functions. This, I need it for

> > > > next steps, but already at this point it improves things as it avoids

> > > > some of the OF parsing, and that's good, isn't it?

> > >

> > > I will take the patches Mark ACKed and send them for v5.2 as

> > > early as it gets in v5.1-rc* cycle.

> >

> > Actually, may I suggest we funnel these through Rafael's tree, unless

> > you are expecting other PSCI changes for v.5.2, which could cause

> > conflicts?

> >

> > The reason is, other PM core changes, to genpd for example, needs to

> > go via Rafael's tree. Those would then potentially block us for

> > applying any other changes to your tree (arm-soc?) for PSCI (as there

> > is dependency) until v5.3.

> >

> > How about if you provides your explicit acks for those PSCI changes

> > your are happy with, then Rafael can pick them?

>

> It is fine we can do that, I would have not sent the patches Mark

> has ACKed to arm-soc till -{rc2/rc3} anyway.


Great!

May I suggest you just reply to the cover-letter and provide the acks
to the relevant patches, then I can then collect the received acks
tags and re-post them to Rafael once rc1 is out.

Kind regards
Uffe

>

> Thanks,

> Lorenzo

>

> > > For this one maybe you can post the changes on top and see what's

> > > the best way forward ?

> > >

> > > I agree that duplicating idle state parsing code across back-ends

> > > is silly - we just want to keep PSCI and kernel data structure

> > > decoupled.

> >

> > Right. Let's continue this discussion later on and move forward here.

> > Make sense.

> >

> > >

> > > Post the code on top and we will find a way forward, OK ?

> >

> > Sure, let me do that.

> >

> > Thanks and kind regards

> > Uffe
Lorenzo Pieralisi March 8, 2019, 1:31 p.m. | #7
On Fri, Mar 08, 2019 at 02:23:26PM +0100, Ulf Hansson wrote:
> On Fri, 8 Mar 2019 at 14:17, Lorenzo Pieralisi

> <lorenzo.pieralisi@arm.com> wrote:

> >

> > On Fri, Mar 08, 2019 at 02:07:51PM +0100, Ulf Hansson wrote:

> > > On Fri, 8 Mar 2019 at 12:49, Lorenzo Pieralisi

> > > <lorenzo.pieralisi@arm.com> wrote:

> > > >

> > > > On Fri, Mar 08, 2019 at 11:36:49AM +0100, Ulf Hansson wrote:

> > > >

> > > > [...]

> > > >

> > > > > Instead, my suggestion is according to what I propose in patch 4 and

> > > > > $subject patch, which means minor adjustments to be able to pass the

> > > > > struct cpuidle_driver * to the init functions. This, I need it for

> > > > > next steps, but already at this point it improves things as it avoids

> > > > > some of the OF parsing, and that's good, isn't it?

> > > >

> > > > I will take the patches Mark ACKed and send them for v5.2 as

> > > > early as it gets in v5.1-rc* cycle.

> > >

> > > Actually, may I suggest we funnel these through Rafael's tree, unless

> > > you are expecting other PSCI changes for v.5.2, which could cause

> > > conflicts?

> > >

> > > The reason is, other PM core changes, to genpd for example, needs to

> > > go via Rafael's tree. Those would then potentially block us for

> > > applying any other changes to your tree (arm-soc?) for PSCI (as there

> > > is dependency) until v5.3.

> > >

> > > How about if you provides your explicit acks for those PSCI changes

> > > your are happy with, then Rafael can pick them?

> >

> > It is fine we can do that, I would have not sent the patches Mark

> > has ACKed to arm-soc till -{rc2/rc3} anyway.

> 

> Great!

> 

> May I suggest you just reply to the cover-letter and provide the acks

> to the relevant patches, then I can then collect the received acks

> tags and re-post them to Rafael once rc1 is out.


Mark ACKed the patches that we consider ready for upstream, tag them
and send them out at -rc1 there is nothing left to do on those.

Lorenzo

Patch

diff --git a/drivers/firmware/psci/psci.c b/drivers/firmware/psci/psci.c
index d50b46a0528f..cbfc936d251c 100644
--- a/drivers/firmware/psci/psci.c
+++ b/drivers/firmware/psci/psci.c
@@ -290,26 +290,20 @@  static int psci_dt_parse_state_node(struct device_node *np, u32 *state)
 static int psci_dt_cpu_init_idle(struct cpuidle_driver *drv,
 			struct device_node *cpu_node, int cpu)
 {
-	int i, ret = 0, count = 0;
+	int i, ret = 0, num_state_nodes = drv->state_count - 1;
 	u32 *psci_states;
 	struct device_node *state_node;
 
-	/* Count idle states */
-	while ((state_node = of_parse_phandle(cpu_node, "cpu-idle-states",
-					      count))) {
-		count++;
-		of_node_put(state_node);
-	}
-
-	if (!count)
-		return -ENODEV;
-
-	psci_states = kcalloc(count, sizeof(*psci_states), GFP_KERNEL);
+	psci_states = kcalloc(num_state_nodes, sizeof(*psci_states),
+			GFP_KERNEL);
 	if (!psci_states)
 		return -ENOMEM;
 
-	for (i = 0; i < count; i++) {
+	for (i = 0; i < num_state_nodes; i++) {
 		state_node = of_parse_phandle(cpu_node, "cpu-idle-states", i);
+		if (!state_node)
+			break;
+
 		ret = psci_dt_parse_state_node(state_node, &psci_states[i]);
 		of_node_put(state_node);
 
@@ -319,6 +313,11 @@  static int psci_dt_cpu_init_idle(struct cpuidle_driver *drv,
 		pr_debug("psci-power-state %#x index %d\n", psci_states[i], i);
 	}
 
+	if (i != num_state_nodes) {
+		ret = -ENODEV;
+		goto free_mem;
+	}
+
 	/* Idle states parsed correctly, initialize per-cpu pointer */
 	per_cpu(psci_power_state, cpu) = psci_states;
 	return 0;