Skip to content

Commit f780d59

Browse files
sklamsklam
authored
Merge pull request #7398 from IntelLabs/parfor_has_shape
Add support for arrayanalysis of tuple args.
2 parents efb0288 + e39f6c9 commit f780d59

3 files changed

Lines changed: 136 additions & 7 deletions

File tree

‎numba/parfors/array_analysis.py‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1325,6 +1325,20 @@ def _analyze_inst(self, label, scope, equiv_set, inst, redefined):
13251325
shape = gvalue
13261326
elif isinstance(gvalue, int):
13271327
shape = (gvalue,)
1328+
elif isinstance(inst.value, ir.Arg):
1329+
if (
1330+
isinstance(typ, types.containers.UniTuple)
1331+
and isinstance(typ.dtype, types.Integer)
1332+
):
1333+
shape = inst.value
1334+
elif (
1335+
isinstance(typ, types.containers.Tuple)
1336+
and all([isinstance(x,
1337+
(types.Integer, types.IntegerLiteral))
1338+
for x in typ.types]
1339+
)
1340+
):
1341+
shape = inst.value
13281342

13291343
if isinstance(shape, ir.Const):
13301344
if isinstance(shape.value, tuple):
@@ -1379,6 +1393,16 @@ def _analyze_inst(self, label, scope, equiv_set, inst, redefined):
13791393
shape = self._gen_shape_call(
13801394
equiv_set, lhs, len(typ), shape, post
13811395
)
1396+
elif (
1397+
isinstance(typ, types.containers.Tuple)
1398+
and all([isinstance(x,
1399+
(types.Integer, types.IntegerLiteral))
1400+
for x in typ.types]
1401+
)
1402+
):
1403+
shape = self._gen_shape_call(
1404+
equiv_set, lhs, len(typ), shape, post
1405+
)
13821406

13831407
""" See the comment on the define() function.
13841408
@@ -3128,11 +3152,16 @@ def _gen_shape_call(self, equiv_set, var, ndims, shape, post):
31283152
# attr call: A_sh_attr = getattr(A, shape)
31293153
if isinstance(shape, ir.Var):
31303154
shape = equiv_set.get_shape(shape)
3155+
31313156
# already a tuple variable that contains size
31323157
if isinstance(shape, ir.Var):
31333158
attr_var = shape
31343159
shape_attr_call = None
31353160
shape = None
3161+
elif isinstance(shape, ir.Arg):
3162+
attr_var = var
3163+
shape_attr_call = None
3164+
shape = None
31363165
else:
31373166
shape_attr_call = ir.Expr.getattr(var, "shape", var.loc)
31383167
attr_var = ir.Var(

‎numba/parfors/parfor.py‎

Lines changed: 25 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1873,10 +1873,18 @@ def _type_getitem(self, args):
18731873
return self.pass_states.typingctx.resolve_function_type(fnty, tuple(args), {})
18741874

18751875

1876-
def _make_index_var(typemap, scope, index_vars, body_block):
1876+
def _make_index_var(typemap, scope, index_vars, body_block, force_tuple=False):
1877+
""" When generating a SetItem call to an array in a parfor, the general
1878+
strategy is to generate a tuple if the array is more than 1 dimension.
1879+
If it is 1 dimensional then you can use a simple variable. This routine
1880+
is also used when converting pndindex to parfor but pndindex requires a
1881+
tuple even if the iteration space is 1 dimensional. The pndindex use of
1882+
this function will use force_tuple to make the output index a tuple even
1883+
if it is one dimensional.
1884+
"""
18771885
ndims = len(index_vars)
18781886
loc = body_block.loc
1879-
if ndims > 1:
1887+
if ndims > 1 or force_tuple:
18801888
tuple_var = ir.Var(scope, mk_unique_var(
18811889
"$parfor_index_tuple_var"), loc)
18821890
typemap[tuple_var.name] = types.containers.UniTuple(
@@ -2446,11 +2454,20 @@ def find_mask_from_size(size_var):
24462454
in_arr, mask_var, mask_typ, mask_indices = result
24472455
else:
24482456
in_arr = args[0]
2449-
size_vars = equiv_set.get_shape(in_arr
2450-
if mask_indices is None else mask_var)
2451-
index_vars, loops = _mk_parfor_loops(
2452-
pass_states.typemap, size_vars, scope, loc,
2453-
)
2457+
assert(isinstance(in_arr, ir.Var))
2458+
in_arr_typ = pass_states.typemap[in_arr.name]
2459+
if isinstance(in_arr_typ, types.Integer):
2460+
index_var = ir.Var(scope, mk_unique_var("parfor_index"), loc)
2461+
pass_states.typemap[index_var.name] = types.uintp
2462+
loops = [LoopNest(index_var, 0, in_arr, 1)]
2463+
index_vars = [index_var]
2464+
else:
2465+
size_vars = equiv_set.get_shape(in_arr
2466+
if mask_indices is None else mask_var)
2467+
index_vars, loops = _mk_parfor_loops(
2468+
pass_states.typemap, size_vars, scope, loc,
2469+
)
2470+
assert(len(loops) > 0)
24542471
orig_index = index_vars
24552472
if mask_indices:
24562473
# replace mask indices if required;
@@ -2462,6 +2479,7 @@ def find_mask_from_size(size_var):
24622479
body_block = ir.Block(scope, loc)
24632480
index_var, index_var_typ = _make_index_var(
24642481
pass_states.typemap, scope, index_vars, body_block,
2482+
force_tuple=True
24652483
)
24662484
body = body_block.body + first_body_block.body
24672485
first_body_block.body = body

‎numba/tests/test_parfors.py‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1874,6 +1874,88 @@ def test_impl(x):
18741874
x = np.ones((2, 2, 2, 2, 2, 15))
18751875
self.check(test_impl, x)
18761876

1877+
def test_tuple_arg(self):
1878+
def test_impl(x, sz):
1879+
for i in numba.pndindex(sz):
1880+
x[i] = 1
1881+
return x
1882+
sz = (10, 5)
1883+
self.check(test_impl, np.empty(sz), sz)
1884+
1885+
def test_tuple_arg_not_whole_array(self):
1886+
def test_impl(x, sz):
1887+
for i in numba.pndindex(sz):
1888+
x[i] = 1
1889+
return x
1890+
sz = (10, 5)
1891+
self.check(test_impl, np.zeros(sz), (10, 3))
1892+
1893+
def test_tuple_for_pndindex(self):
1894+
def test_impl(x):
1895+
sz = (10, 5)
1896+
for i in numba.pndindex(sz):
1897+
x[i] = 1
1898+
return x
1899+
sz = (10, 5)
1900+
self.check(test_impl, np.zeros(sz))
1901+
1902+
def test_tuple_arg_literal(self):
1903+
def test_impl(x, first):
1904+
sz = (first, 5)
1905+
for i in numba.pndindex(sz):
1906+
x[i] = 1
1907+
return x
1908+
sz = (10, 5)
1909+
self.check(test_impl, np.zeros(sz), 10)
1910+
1911+
def test_tuple_of_literal_nonliteral(self):
1912+
# This test has to be done manually as the self.check uses
1913+
# compile_isolated and one function cannot "see" the other
1914+
1915+
def test_impl(x, sz):
1916+
for i in numba.pndindex(sz):
1917+
x[i] = 1
1918+
return x
1919+
1920+
def call(x, fn):
1921+
return fn(x, (10, 3)) # Only want to iterate to the 3rd
1922+
1923+
get_input = lambda: np.zeros((10, 10))
1924+
expected = call(get_input(), test_impl)
1925+
1926+
def check(dec):
1927+
f1 = dec(test_impl)
1928+
f2 = njit(call) # no parallel semantics in the caller
1929+
got = f2(get_input(), f1)
1930+
self.assertPreciseEqual(expected, got)
1931+
1932+
for d in (njit, njit(parallel=True)):
1933+
check(d)
1934+
1935+
def test_tuple_arg_1d(self):
1936+
def test_impl(x, sz):
1937+
for i in numba.pndindex(sz):
1938+
x[i] = 1
1939+
return x
1940+
sz = (10,)
1941+
self.check(test_impl, np.zeros(sz), sz)
1942+
1943+
def test_tuple_arg_1d_literal(self):
1944+
def test_impl(x):
1945+
sz = (10,)
1946+
for i in numba.pndindex(sz):
1947+
x[i] = 1
1948+
return x
1949+
sz = (10,)
1950+
self.check(test_impl, np.zeros(sz))
1951+
1952+
def test_int_arg_pndindex(self):
1953+
def test_impl(x, sz):
1954+
for i in numba.pndindex(sz):
1955+
x[i] = 1
1956+
return x
1957+
self.check(test_impl, np.zeros((10, 10)), 3)
1958+
18771959

18781960
@skip_parfors_unsupported
18791961
class TestParforsLeaks(MemoryLeakMixin, TestParforsBase):

0 commit comments

Comments
 (0)