diff mbox series

[v6,6/6] powerpc/pseries: Consolidate form1 distance initialization into a helper

Message ID 20210727100311.310969-7-aneesh.kumar@linux.ibm.com (mailing list archive)
State Superseded
Headers show
Series Add support for FORM2 associativity | expand
Related show

Checks

Context Check Description
snowpatch_ozlabs/github-powerpc_selftests success Successfully ran 8 jobs.
snowpatch_ozlabs/github-powerpc_sparse success Successfully ran 4 jobs.
snowpatch_ozlabs/github-powerpc_ppctests success Successfully ran 8 jobs.
snowpatch_ozlabs/github-powerpc_clang success Successfully ran 8 jobs.
snowpatch_ozlabs/github-powerpc_kernel_qemu success Successfully ran 25 jobs.

Commit Message

Aneesh Kumar K V July 27, 2021, 10:03 a.m. UTC
Currently, we duplicate parsing code for ibm,associativity and
ibm,associativity-lookup-arrays in the kernel. The associativity array provided
by these device tree properties are very similar and hence can use
a helper to parse the node id and numa distance details.

Signed-off-by: Aneesh Kumar K.V <aneesh.kumar@linux.ibm.com>
---
 arch/powerpc/mm/numa.c | 83 ++++++++++++++++++++++++++----------------
 1 file changed, 51 insertions(+), 32 deletions(-)

Comments

David Gibson Aug. 6, 2021, 6:47 a.m. UTC | #1
On Tue, Jul 27, 2021 at 03:33:11PM +0530, Aneesh Kumar K.V wrote:
> Currently, we duplicate parsing code for ibm,associativity and
> ibm,associativity-lookup-arrays in the kernel. The associativity array provided
> by these device tree properties are very similar and hence can use
> a helper to parse the node id and numa distance details.

Oh... sorry.. comments on the earlier patch were from before I read
and saw you adjusted things here.

> 
> Signed-off-by: Aneesh Kumar K.V <aneesh.kumar@linux.ibm.com>
> ---
>  arch/powerpc/mm/numa.c | 83 ++++++++++++++++++++++++++----------------
>  1 file changed, 51 insertions(+), 32 deletions(-)
> 
> diff --git a/arch/powerpc/mm/numa.c b/arch/powerpc/mm/numa.c
> index fffb3c40f595..7506251e17f2 100644
> --- a/arch/powerpc/mm/numa.c
> +++ b/arch/powerpc/mm/numa.c
> @@ -171,19 +171,19 @@ static void unmap_cpu_from_node(unsigned long cpu)
>  }
>  #endif /* CONFIG_HOTPLUG_CPU || CONFIG_PPC_SPLPAR */
>  
> -/*
> - * Returns nid in the range [0..nr_node_ids], or -1 if no useful NUMA
> - * info is found.
> - */
> -static int associativity_to_nid(const __be32 *associativity)
> +static int __associativity_to_nid(const __be32 *associativity,
> +				  int max_array_sz)
>  {
>  	int nid = NUMA_NO_NODE;
> +	/*
> +	 * primary_domain_index is 1 based array index.
> +	 */
> +	int index = primary_domain_index  - 1;
>  
> -	if (!numa_enabled)
> +	if (!numa_enabled || index >= max_array_sz)
>  		goto out;

You don't need a goto, you can just return NUMA_NO_NODE.

>  
> -	if (of_read_number(associativity, 1) >= primary_domain_index)
> -		nid = of_read_number(&associativity[primary_domain_index], 1);
> +	nid = of_read_number(&associativity[index], 1);
>  
>  	/* POWER4 LPAR uses 0xffff as invalid node */
>  	if (nid == 0xffff || nid >= nr_node_ids)
> @@ -191,6 +191,17 @@ static int associativity_to_nid(const __be32 *associativity)
>  out:
>  	return nid;
>  }
> +/*
> + * Returns nid in the range [0..nr_node_ids], or -1 if no useful NUMA
> + * info is found.
> + */
> +static int associativity_to_nid(const __be32 *associativity)
> +{
> +	int array_sz = of_read_number(associativity, 1);
> +
> +	/* Skip the first element in the associativity array */
> +	return __associativity_to_nid((associativity + 1), array_sz);
> +}
>  
>  static int __cpu_form2_relative_distance(__be32 *cpu1_assoc, __be32 *cpu2_assoc)
>  {
> @@ -295,24 +306,41 @@ int of_node_to_nid(struct device_node *device)
>  }
>  EXPORT_SYMBOL(of_node_to_nid);
>  
> -static void __initialize_form1_numa_distance(const __be32 *associativity)
> +static void ___initialize_form1_numa_distance(const __be32 *associativity,
> +					     int max_array_sz)
>  {
>  	int i, nid;
>  
>  	if (affinity_form != FORM1_AFFINITY)
>  		return;
>  
> -	nid = associativity_to_nid(associativity);
> +	nid = __associativity_to_nid(associativity, max_array_sz);
>  	if (nid != NUMA_NO_NODE) {
>  		for (i = 0; i < distance_ref_points_depth; i++) {
>  			const __be32 *entry;
> +			int index = be32_to_cpu(distance_ref_points[i]) - 1;
> +
> +			/*
> +			 * broken hierarchy, return with broken distance table

WARN_ON, maybe?

> +			 */
> +			if (index >= max_array_sz)
> +				return;
>  
> -			entry = &associativity[be32_to_cpu(distance_ref_points[i])];
> +			entry = &associativity[index];
>  			distance_lookup_table[nid][i] = of_read_number(entry, 1);
>  		}
>  	}
>  }
>  
> +static void __initialize_form1_numa_distance(const __be32 *associativity)

Do you actually use this in-between wrapper?

> +{
> +	int array_sz;
> +
> +	array_sz = of_read_number(associativity, 1);
> +	/* Skip the first element in the associativity array */
> +	___initialize_form1_numa_distance(associativity + 1, array_sz);
> +}
> +
>  static void initialize_form1_numa_distance(struct device_node *node)
>  {
>  	const __be32 *associativity;
> @@ -586,27 +614,18 @@ static int get_nid_and_numa_distance(struct drmem_lmb *lmb)
>  
>  	if (primary_domain_index <= aa.array_sz &&
>  	    !(lmb->flags & DRCONF_MEM_AI_INVALID) && lmb->aa_index < aa.n_arrays) {
> -		index = lmb->aa_index * aa.array_sz + primary_domain_index - 1;
> -		nid = of_read_number(&aa.arrays[index], 1);
> +		const __be32 *associativity;
>  
> -		if (nid == 0xffff || nid >= nr_node_ids)
> -			nid = default_nid;
> +		index = lmb->aa_index * aa.array_sz;
> +		associativity = &aa.arrays[index];
> +		nid = __associativity_to_nid(associativity, aa.array_sz);
>  		if (nid > 0 && affinity_form == FORM1_AFFINITY) {
> -			int i;
> -			const __be32 *associativity;
> -
> -			index = lmb->aa_index * aa.array_sz;
> -			associativity = &aa.arrays[index];
>  			/*
> -			 * lookup array associativity entries have different format
> -			 * There is no length of the array as the first element.
> +			 * lookup array associativity entries have
> +			 * no length of the array as the first element.
>  			 */
> -			for (i = 0; i < distance_ref_points_depth; i++) {
> -				const __be32 *entry;
> -
> -				entry = &associativity[be32_to_cpu(distance_ref_points[i]) - 1];
> -				distance_lookup_table[nid][i] = of_read_number(entry, 1);
> -			}
> +			___initialize_form1_numa_distance(associativity,
> +							  aa.array_sz);

Better, thanks.

>  		}
>  	}
>  	return nid;
> @@ -632,11 +651,11 @@ int of_drconf_to_nid_single(struct drmem_lmb *lmb)
>  
>  	if (primary_domain_index <= aa.array_sz &&
>  	    !(lmb->flags & DRCONF_MEM_AI_INVALID) && lmb->aa_index < aa.n_arrays) {
> -		index = lmb->aa_index * aa.array_sz + primary_domain_index - 1;
> -		nid = of_read_number(&aa.arrays[index], 1);
> +		const __be32 *associativity;
>  
> -		if (nid == 0xffff || nid >= nr_node_ids)
> -			nid = default_nid;
> +		index = lmb->aa_index * aa.array_sz;
> +		associativity = &aa.arrays[index];
> +		nid = __associativity_to_nid(associativity, aa.array_sz);
>  	}
>  	return nid;
>  }
Aneesh Kumar K V Aug. 6, 2021, 4:23 p.m. UTC | #2
On 8/6/21 12:17 PM, David Gibson wrote:
> On Tue, Jul 27, 2021 at 03:33:11PM +0530, Aneesh Kumar K.V wrote:
>> Currently, we duplicate parsing code for ibm,associativity and
>> ibm,associativity-lookup-arrays in the kernel. The associativity array provided
>> by these device tree properties are very similar and hence can use
>> a helper to parse the node id and numa distance details.
> 
> Oh... sorry.. comments on the earlier patch were from before I read
> and saw you adjusted things here.
> 
>>
>> Signed-off-by: Aneesh Kumar K.V <aneesh.kumar@linux.ibm.com>
>> ---
>>   arch/powerpc/mm/numa.c | 83 ++++++++++++++++++++++++++----------------
>>   1 file changed, 51 insertions(+), 32 deletions(-)
>>
>> diff --git a/arch/powerpc/mm/numa.c b/arch/powerpc/mm/numa.c
>> index fffb3c40f595..7506251e17f2 100644
>> --- a/arch/powerpc/mm/numa.c
>> +++ b/arch/powerpc/mm/numa.c
>> @@ -171,19 +171,19 @@ static void unmap_cpu_from_node(unsigned long cpu)
>>   }
>>   #endif /* CONFIG_HOTPLUG_CPU || CONFIG_PPC_SPLPAR */
>>   
>> -/*
>> - * Returns nid in the range [0..nr_node_ids], or -1 if no useful NUMA
>> - * info is found.
>> - */
>> -static int associativity_to_nid(const __be32 *associativity)
>> +static int __associativity_to_nid(const __be32 *associativity,
>> +				  int max_array_sz)
>>   {
>>   	int nid = NUMA_NO_NODE;
>> +	/*
>> +	 * primary_domain_index is 1 based array index.
>> +	 */
>> +	int index = primary_domain_index  - 1;
>>   
>> -	if (!numa_enabled)
>> +	if (!numa_enabled || index >= max_array_sz)
>>   		goto out;
> 
> You don't need a goto, you can just return NUMA_NO_NODE.

updated

> 
>>   
>> -	if (of_read_number(associativity, 1) >= primary_domain_index)
>> -		nid = of_read_number(&associativity[primary_domain_index], 1);
>> +	nid = of_read_number(&associativity[index], 1);
>>   
>>   	/* POWER4 LPAR uses 0xffff as invalid node */
>>   	if (nid == 0xffff || nid >= nr_node_ids)
>> @@ -191,6 +191,17 @@ static int associativity_to_nid(const __be32 *associativity)
>>   out:
>>   	return nid;
>>   }
>> +/*
>> + * Returns nid in the range [0..nr_node_ids], or -1 if no useful NUMA
>> + * info is found.
>> + */
>> +static int associativity_to_nid(const __be32 *associativity)
>> +{
>> +	int array_sz = of_read_number(associativity, 1);
>> +
>> +	/* Skip the first element in the associativity array */
>> +	return __associativity_to_nid((associativity + 1), array_sz);
>> +}
>>   
>>   static int __cpu_form2_relative_distance(__be32 *cpu1_assoc, __be32 *cpu2_assoc)
>>   {
>> @@ -295,24 +306,41 @@ int of_node_to_nid(struct device_node *device)
>>   }
>>   EXPORT_SYMBOL(of_node_to_nid);
>>   
>> -static void __initialize_form1_numa_distance(const __be32 *associativity)
>> +static void ___initialize_form1_numa_distance(const __be32 *associativity,
>> +					     int max_array_sz)
>>   {
>>   	int i, nid;
>>   
>>   	if (affinity_form != FORM1_AFFINITY)
>>   		return;
>>   
>> -	nid = associativity_to_nid(associativity);
>> +	nid = __associativity_to_nid(associativity, max_array_sz);
>>   	if (nid != NUMA_NO_NODE) {
>>   		for (i = 0; i < distance_ref_points_depth; i++) {
>>   			const __be32 *entry;
>> +			int index = be32_to_cpu(distance_ref_points[i]) - 1;
>> +
>> +			/*
>> +			 * broken hierarchy, return with broken distance table
> 
> WARN_ON, maybe?


updated

> 
>> +			 */
>> +			if (index >= max_array_sz)
>> +				return;
>>   
>> -			entry = &associativity[be32_to_cpu(distance_ref_points[i])];
>> +			entry = &associativity[index];
>>   			distance_lookup_table[nid][i] = of_read_number(entry, 1);
>>   		}
>>   	}
>>   }
>>   
>> +static void __initialize_form1_numa_distance(const __be32 *associativity)
> 
> Do you actually use this in-between wrapper?

yes used in

static void initialize_form1_numa_distance(struct device_node *node)
{
	const __be32 *associativity;

	associativity = of_get_associativity(node);
	if (!associativity)
		return;

	__initialize_form1_numa_distance(associativity);
}



> 
>> +{
>> +	int array_sz;
>> +
>> +	array_sz = of_read_number(associativity, 1);
>> +	/* Skip the first element in the associativity array */
>> +	___initialize_form1_numa_distance(associativity + 1, array_sz);
>> +}
>> +
>>   static void initialize_form1_numa_distance(struct device_node *node)
>>   {
>>   	const __be32 *associativity;
>> @@ -586,27 +614,18 @@ static int get_nid_and_numa_distance(struct drmem_lmb *lmb)
>>   
>>   	if (primary_domain_index <= aa.array_sz &&
>>   	    !(lmb->flags & DRCONF_MEM_AI_INVALID) && lmb->aa_index < aa.n_arrays) {
>> -		index = lmb->aa_index * aa.array_sz + primary_domain_index - 1;
>> -		nid = of_read_number(&aa.arrays[index], 1);
>> +		const __be32 *associativity;
>>   
>> -		if (nid == 0xffff || nid >= nr_node_ids)
>> -			nid = default_nid;
>> +		index = lmb->aa_index * aa.array_sz;
>> +		associativity = &aa.arrays[index];
>> +		nid = __associativity_to_nid(associativity, aa.array_sz);
>>   		if (nid > 0 && affinity_form == FORM1_AFFINITY) {
>> -			int i;
>> -			const __be32 *associativity;
>> -
>> -			index = lmb->aa_index * aa.array_sz;
>> -			associativity = &aa.arrays[index];
>>   			/*
>> -			 * lookup array associativity entries have different format
>> -			 * There is no length of the array as the first element.
>> +			 * lookup array associativity entries have
>> +			 * no length of the array as the first element.
>>   			 */
>> -			for (i = 0; i < distance_ref_points_depth; i++) {
>> -				const __be32 *entry;
>> -
>> -				entry = &associativity[be32_to_cpu(distance_ref_points[i]) - 1];
>> -				distance_lookup_table[nid][i] = of_read_number(entry, 1);
>> -			}
>> +			___initialize_form1_numa_distance(associativity,
>> +							  aa.array_sz);
> 
> Better, thanks.
> 
>>   		}
>>   	}
>>   	return nid;
>> @@ -632,11 +651,11 @@ int of_drconf_to_nid_single(struct drmem_lmb *lmb)
>>   
>>   	if (primary_domain_index <= aa.array_sz &&
>>   	    !(lmb->flags & DRCONF_MEM_AI_INVALID) && lmb->aa_index < aa.n_arrays) {
>> -		index = lmb->aa_index * aa.array_sz + primary_domain_index - 1;
>> -		nid = of_read_number(&aa.arrays[index], 1);
>> +		const __be32 *associativity;
>>   
>> -		if (nid == 0xffff || nid >= nr_node_ids)
>> -			nid = default_nid;
>> +		index = lmb->aa_index * aa.array_sz;
>> +		associativity = &aa.arrays[index];
>> +		nid = __associativity_to_nid(associativity, aa.array_sz);
>>   	}
>>   	return nid;
>>   }
> 


-aneesh
David Gibson Aug. 9, 2021, 3:27 a.m. UTC | #3
On Fri, Aug 06, 2021 at 09:53:59PM +0530, Aneesh Kumar K.V wrote:
> On 8/6/21 12:17 PM, David Gibson wrote:
> > On Tue, Jul 27, 2021 at 03:33:11PM +0530, Aneesh Kumar K.V wrote:
> > > Currently, we duplicate parsing code for ibm,associativity and
> > > ibm,associativity-lookup-arrays in the kernel. The associativity array provided
> > > by these device tree properties are very similar and hence can use
> > > a helper to parse the node id and numa distance details.
> > 
> > Oh... sorry.. comments on the earlier patch were from before I read
> > and saw you adjusted things here.
> > 
> > > 
> > > Signed-off-by: Aneesh Kumar K.V <aneesh.kumar@linux.ibm.com>
> > > ---
> > >   arch/powerpc/mm/numa.c | 83 ++++++++++++++++++++++++++----------------
> > >   1 file changed, 51 insertions(+), 32 deletions(-)
> > > 
> > > diff --git a/arch/powerpc/mm/numa.c b/arch/powerpc/mm/numa.c
> > > index fffb3c40f595..7506251e17f2 100644
> > > --- a/arch/powerpc/mm/numa.c
> > > +++ b/arch/powerpc/mm/numa.c
> > > @@ -171,19 +171,19 @@ static void unmap_cpu_from_node(unsigned long cpu)
> > >   }
> > >   #endif /* CONFIG_HOTPLUG_CPU || CONFIG_PPC_SPLPAR */
> > > -/*
> > > - * Returns nid in the range [0..nr_node_ids], or -1 if no useful NUMA
> > > - * info is found.
> > > - */
> > > -static int associativity_to_nid(const __be32 *associativity)
> > > +static int __associativity_to_nid(const __be32 *associativity,
> > > +				  int max_array_sz)
> > >   {
> > >   	int nid = NUMA_NO_NODE;
> > > +	/*
> > > +	 * primary_domain_index is 1 based array index.
> > > +	 */
> > > +	int index = primary_domain_index  - 1;
> > > -	if (!numa_enabled)
> > > +	if (!numa_enabled || index >= max_array_sz)
> > >   		goto out;
> > 
> > You don't need a goto, you can just return NUMA_NO_NODE.
> 
> updated
> 
> > 
> > > -	if (of_read_number(associativity, 1) >= primary_domain_index)
> > > -		nid = of_read_number(&associativity[primary_domain_index], 1);
> > > +	nid = of_read_number(&associativity[index], 1);
> > >   	/* POWER4 LPAR uses 0xffff as invalid node */
> > >   	if (nid == 0xffff || nid >= nr_node_ids)
> > > @@ -191,6 +191,17 @@ static int associativity_to_nid(const __be32 *associativity)
> > >   out:
> > >   	return nid;
> > >   }
> > > +/*
> > > + * Returns nid in the range [0..nr_node_ids], or -1 if no useful NUMA
> > > + * info is found.
> > > + */
> > > +static int associativity_to_nid(const __be32 *associativity)
> > > +{
> > > +	int array_sz = of_read_number(associativity, 1);
> > > +
> > > +	/* Skip the first element in the associativity array */
> > > +	return __associativity_to_nid((associativity + 1), array_sz);
> > > +}
> > >   static int __cpu_form2_relative_distance(__be32 *cpu1_assoc, __be32 *cpu2_assoc)
> > >   {
> > > @@ -295,24 +306,41 @@ int of_node_to_nid(struct device_node *device)
> > >   }
> > >   EXPORT_SYMBOL(of_node_to_nid);
> > > -static void __initialize_form1_numa_distance(const __be32 *associativity)
> > > +static void ___initialize_form1_numa_distance(const __be32 *associativity,
> > > +					     int max_array_sz)
> > >   {
> > >   	int i, nid;
> > >   	if (affinity_form != FORM1_AFFINITY)
> > >   		return;
> > > -	nid = associativity_to_nid(associativity);
> > > +	nid = __associativity_to_nid(associativity, max_array_sz);
> > >   	if (nid != NUMA_NO_NODE) {
> > >   		for (i = 0; i < distance_ref_points_depth; i++) {
> > >   			const __be32 *entry;
> > > +			int index = be32_to_cpu(distance_ref_points[i]) - 1;
> > > +
> > > +			/*
> > > +			 * broken hierarchy, return with broken distance table
> > 
> > WARN_ON, maybe?
> 
> 
> updated
> 
> > 
> > > +			 */
> > > +			if (index >= max_array_sz)
> > > +				return;
> > > -			entry = &associativity[be32_to_cpu(distance_ref_points[i])];
> > > +			entry = &associativity[index];
> > >   			distance_lookup_table[nid][i] = of_read_number(entry, 1);
> > >   		}
> > >   	}
> > >   }
> > > +static void __initialize_form1_numa_distance(const __be32 *associativity)
> > 
> > Do you actually use this in-between wrapper?
> 
> yes used in
> 
> static void initialize_form1_numa_distance(struct device_node *node)
> {
> 	const __be32 *associativity;
> 
> 	associativity = of_get_associativity(node);
> 	if (!associativity)
> 		return;
> 
> 	__initialize_form1_numa_distance(associativity);

Right, but I mean in more than one place.  If this is the only user,
you might as well expand it inline here.

> }
> 
> 
> 
> > 
> > > +{
> > > +	int array_sz;
> > > +
> > > +	array_sz = of_read_number(associativity, 1);
> > > +	/* Skip the first element in the associativity array */
> > > +	___initialize_form1_numa_distance(associativity + 1, array_sz);
> > > +}
> > > +
> > >   static void initialize_form1_numa_distance(struct device_node *node)
> > >   {
> > >   	const __be32 *associativity;
> > > @@ -586,27 +614,18 @@ static int get_nid_and_numa_distance(struct drmem_lmb *lmb)
> > >   	if (primary_domain_index <= aa.array_sz &&
> > >   	    !(lmb->flags & DRCONF_MEM_AI_INVALID) && lmb->aa_index < aa.n_arrays) {
> > > -		index = lmb->aa_index * aa.array_sz + primary_domain_index - 1;
> > > -		nid = of_read_number(&aa.arrays[index], 1);
> > > +		const __be32 *associativity;
> > > -		if (nid == 0xffff || nid >= nr_node_ids)
> > > -			nid = default_nid;
> > > +		index = lmb->aa_index * aa.array_sz;
> > > +		associativity = &aa.arrays[index];
> > > +		nid = __associativity_to_nid(associativity, aa.array_sz);
> > >   		if (nid > 0 && affinity_form == FORM1_AFFINITY) {
> > > -			int i;
> > > -			const __be32 *associativity;
> > > -
> > > -			index = lmb->aa_index * aa.array_sz;
> > > -			associativity = &aa.arrays[index];
> > >   			/*
> > > -			 * lookup array associativity entries have different format
> > > -			 * There is no length of the array as the first element.
> > > +			 * lookup array associativity entries have
> > > +			 * no length of the array as the first element.
> > >   			 */
> > > -			for (i = 0; i < distance_ref_points_depth; i++) {
> > > -				const __be32 *entry;
> > > -
> > > -				entry = &associativity[be32_to_cpu(distance_ref_points[i]) - 1];
> > > -				distance_lookup_table[nid][i] = of_read_number(entry, 1);
> > > -			}
> > > +			___initialize_form1_numa_distance(associativity,
> > > +							  aa.array_sz);
> > 
> > Better, thanks.
> > 
> > >   		}
> > >   	}
> > >   	return nid;
> > > @@ -632,11 +651,11 @@ int of_drconf_to_nid_single(struct drmem_lmb *lmb)
> > >   	if (primary_domain_index <= aa.array_sz &&
> > >   	    !(lmb->flags & DRCONF_MEM_AI_INVALID) && lmb->aa_index < aa.n_arrays) {
> > > -		index = lmb->aa_index * aa.array_sz + primary_domain_index - 1;
> > > -		nid = of_read_number(&aa.arrays[index], 1);
> > > +		const __be32 *associativity;
> > > -		if (nid == 0xffff || nid >= nr_node_ids)
> > > -			nid = default_nid;
> > > +		index = lmb->aa_index * aa.array_sz;
> > > +		associativity = &aa.arrays[index];
> > > +		nid = __associativity_to_nid(associativity, aa.array_sz);
> > >   	}
> > >   	return nid;
> > >   }
> > 
> 
> 
> -aneesh
>
diff mbox series

Patch

diff --git a/arch/powerpc/mm/numa.c b/arch/powerpc/mm/numa.c
index fffb3c40f595..7506251e17f2 100644
--- a/arch/powerpc/mm/numa.c
+++ b/arch/powerpc/mm/numa.c
@@ -171,19 +171,19 @@  static void unmap_cpu_from_node(unsigned long cpu)
 }
 #endif /* CONFIG_HOTPLUG_CPU || CONFIG_PPC_SPLPAR */
 
-/*
- * Returns nid in the range [0..nr_node_ids], or -1 if no useful NUMA
- * info is found.
- */
-static int associativity_to_nid(const __be32 *associativity)
+static int __associativity_to_nid(const __be32 *associativity,
+				  int max_array_sz)
 {
 	int nid = NUMA_NO_NODE;
+	/*
+	 * primary_domain_index is 1 based array index.
+	 */
+	int index = primary_domain_index  - 1;
 
-	if (!numa_enabled)
+	if (!numa_enabled || index >= max_array_sz)
 		goto out;
 
-	if (of_read_number(associativity, 1) >= primary_domain_index)
-		nid = of_read_number(&associativity[primary_domain_index], 1);
+	nid = of_read_number(&associativity[index], 1);
 
 	/* POWER4 LPAR uses 0xffff as invalid node */
 	if (nid == 0xffff || nid >= nr_node_ids)
@@ -191,6 +191,17 @@  static int associativity_to_nid(const __be32 *associativity)
 out:
 	return nid;
 }
+/*
+ * Returns nid in the range [0..nr_node_ids], or -1 if no useful NUMA
+ * info is found.
+ */
+static int associativity_to_nid(const __be32 *associativity)
+{
+	int array_sz = of_read_number(associativity, 1);
+
+	/* Skip the first element in the associativity array */
+	return __associativity_to_nid((associativity + 1), array_sz);
+}
 
 static int __cpu_form2_relative_distance(__be32 *cpu1_assoc, __be32 *cpu2_assoc)
 {
@@ -295,24 +306,41 @@  int of_node_to_nid(struct device_node *device)
 }
 EXPORT_SYMBOL(of_node_to_nid);
 
-static void __initialize_form1_numa_distance(const __be32 *associativity)
+static void ___initialize_form1_numa_distance(const __be32 *associativity,
+					     int max_array_sz)
 {
 	int i, nid;
 
 	if (affinity_form != FORM1_AFFINITY)
 		return;
 
-	nid = associativity_to_nid(associativity);
+	nid = __associativity_to_nid(associativity, max_array_sz);
 	if (nid != NUMA_NO_NODE) {
 		for (i = 0; i < distance_ref_points_depth; i++) {
 			const __be32 *entry;
+			int index = be32_to_cpu(distance_ref_points[i]) - 1;
+
+			/*
+			 * broken hierarchy, return with broken distance table
+			 */
+			if (index >= max_array_sz)
+				return;
 
-			entry = &associativity[be32_to_cpu(distance_ref_points[i])];
+			entry = &associativity[index];
 			distance_lookup_table[nid][i] = of_read_number(entry, 1);
 		}
 	}
 }
 
+static void __initialize_form1_numa_distance(const __be32 *associativity)
+{
+	int array_sz;
+
+	array_sz = of_read_number(associativity, 1);
+	/* Skip the first element in the associativity array */
+	___initialize_form1_numa_distance(associativity + 1, array_sz);
+}
+
 static void initialize_form1_numa_distance(struct device_node *node)
 {
 	const __be32 *associativity;
@@ -586,27 +614,18 @@  static int get_nid_and_numa_distance(struct drmem_lmb *lmb)
 
 	if (primary_domain_index <= aa.array_sz &&
 	    !(lmb->flags & DRCONF_MEM_AI_INVALID) && lmb->aa_index < aa.n_arrays) {
-		index = lmb->aa_index * aa.array_sz + primary_domain_index - 1;
-		nid = of_read_number(&aa.arrays[index], 1);
+		const __be32 *associativity;
 
-		if (nid == 0xffff || nid >= nr_node_ids)
-			nid = default_nid;
+		index = lmb->aa_index * aa.array_sz;
+		associativity = &aa.arrays[index];
+		nid = __associativity_to_nid(associativity, aa.array_sz);
 		if (nid > 0 && affinity_form == FORM1_AFFINITY) {
-			int i;
-			const __be32 *associativity;
-
-			index = lmb->aa_index * aa.array_sz;
-			associativity = &aa.arrays[index];
 			/*
-			 * lookup array associativity entries have different format
-			 * There is no length of the array as the first element.
+			 * lookup array associativity entries have
+			 * no length of the array as the first element.
 			 */
-			for (i = 0; i < distance_ref_points_depth; i++) {
-				const __be32 *entry;
-
-				entry = &associativity[be32_to_cpu(distance_ref_points[i]) - 1];
-				distance_lookup_table[nid][i] = of_read_number(entry, 1);
-			}
+			___initialize_form1_numa_distance(associativity,
+							  aa.array_sz);
 		}
 	}
 	return nid;
@@ -632,11 +651,11 @@  int of_drconf_to_nid_single(struct drmem_lmb *lmb)
 
 	if (primary_domain_index <= aa.array_sz &&
 	    !(lmb->flags & DRCONF_MEM_AI_INVALID) && lmb->aa_index < aa.n_arrays) {
-		index = lmb->aa_index * aa.array_sz + primary_domain_index - 1;
-		nid = of_read_number(&aa.arrays[index], 1);
+		const __be32 *associativity;
 
-		if (nid == 0xffff || nid >= nr_node_ids)
-			nid = default_nid;
+		index = lmb->aa_index * aa.array_sz;
+		associativity = &aa.arrays[index];
+		nid = __associativity_to_nid(associativity, aa.array_sz);
 	}
 	return nid;
 }