Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions numba/core/datamodel/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -1077,6 +1077,7 @@ def __init__(self, dmm, fe_type):
super(UniTupleIter, self).__init__(dmm, fe_type, members)


@register_default(types.misc.SliceLiteral)
@register_default(types.SliceType)
class SliceModel(StructModel):
def __init__(self, dmm, fe_type):
Expand Down
5 changes: 5 additions & 0 deletions numba/core/types/misc.py
Original file line number Diff line number Diff line change
Expand Up @@ -334,6 +334,11 @@ def __init__(self, value):
members = 2 if value.step is None else 3
SliceType.__init__(self, name=name, members=members)

@property
def key(self):
sl = self.literal_value
return sl.start, sl.stop, sl.step


Literal.ctor_map[slice] = SliceLiteral

Expand Down
81 changes: 74 additions & 7 deletions numba/cpython/slicing.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,11 @@
from itertools import zip_longest

from llvmlite import ir

from numba.core import types, typing, utils, cgutils
from numba.core.imputils import (lower_builtin, lower_getattr,
iternext_impl, impl_ret_borrowed,
impl_ret_new_ref, impl_ret_untracked)
from numba.core import cgutils, types, typing, utils
from numba.core.imputils import (impl_ret_borrowed, impl_ret_new_ref,
impl_ret_untracked, iternext_impl,
lower_builtin, lower_cast, lower_constant,
lower_getattr)


def fix_index(builder, idx, size):
Expand Down Expand Up @@ -150,8 +150,13 @@ def get_defaults(context):

@lower_builtin(slice, types.VarArg(types.Any))
def slice_constructor_impl(context, builder, sig, args):
default_start_pos, default_start_neg, default_stop_pos, default_stop_neg, default_step = \
[context.get_constant(types.intp, x) for x in get_defaults(context)]
(
default_start_pos,
default_start_neg,
default_stop_pos,
default_stop_neg,
default_step,
) = [context.get_constant(types.intp, x) for x in get_defaults(context)]

slice_args = [None] * 3

Expand Down Expand Up @@ -233,3 +238,65 @@ def slice_indices(context, builder, sig, args):
sig.return_type,
(sli.start, sli.stop, sli.step)
)


def make_slice_from_constant(context, builder, ty, pyval):
sli = context.make_helper(builder, ty)
lty = context.get_value_type(types.intp)

(
default_start_pos,
default_start_neg,
default_stop_pos,
default_stop_neg,
default_step,
) = [context.get_constant(types.intp, x) for x in get_defaults(context)]

step = pyval.step
if step is None:
step_is_neg = False
step = default_step
else:
step_is_neg = step < 0
step = lty(step)

start = pyval.start
if start is None:
if step_is_neg:
start = default_start_neg
else:
start = default_start_pos
else:
start = lty(start)

stop = pyval.stop
if stop is None:
if step_is_neg:
stop = default_stop_neg
else:
stop = default_stop_pos
else:
stop = lty(stop)

sli.start = start
sli.stop = stop
sli.step = step

return sli._getvalue()


@lower_constant(types.SliceType)
def constant_slice(context, builder, ty, pyval):
if isinstance(ty, types.Literal):
typ = ty.literal_type
else:
typ = ty
Comment on lines +292 to +293

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I didn't notice this before... but is this branch dead or maybe it should raise? Is it possible to get a slice type that is constant and not literal?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

From OOB discussion with @sklam, it shouldn't be possible to hit this from user defined input but synthetic IR might.


return make_slice_from_constant(context, builder, typ, pyval)


@lower_cast(types.misc.SliceLiteral, types.SliceType)
def cast_from_literal(context, builder, fromty, toty, val):
return make_slice_from_constant(
context, builder, toty, fromty.literal_value,
)
26 changes: 26 additions & 0 deletions numba/tests/test_parfors.py
Original file line number Diff line number Diff line change
Expand Up @@ -3433,6 +3433,32 @@ def test_impl():

self.check(test_impl)

@skip_parfors_unsupported
def test_parfor_array_access_lower_slice(self):
for ts in [slice(1, 3, None), slice(2, None, None), slice(None, 2, -1),
slice(None, None, None), slice(None, None, -2)]:

def test_impl(n):
X = np.arange(n * 4).reshape((n, 4))
y = 0
for i in numba.prange(n):
y += X[i, ts].sum()
return y

n = 10
self.check(test_impl, n)

X = np.arange(n * 4).reshape((n, 4))

def test_impl(X):
y = 0
for i in numba.prange(X.shape[0]):
y += X[i, ts].sum()
return y

self.check(test_impl, X)



class TestParforsOptions(TestParforsBase):

Expand Down
24 changes: 23 additions & 1 deletion numba/tests/test_slices.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
import numpy as np

from numba import jit, typeof, TypingError
from numba.core import utils
from numba.core import utils, types
from numba.tests.support import TestCase, MemoryLeakMixin
from numba.core.types.functions import _header_lead
import unittest
Expand Down Expand Up @@ -185,6 +185,28 @@ def test_slice_indices_examples(self):
str(e.exception)
)

def test_slice_from_constant(self):
test_tuple = (1, 2, 3, 4)

for ts in itertools.product(
[None, 1, 2, 3], [None, 1, 2, 3], [None, 1, 2, -1, -2]
):
ts = slice(*ts)

@jit(nopython=True)
def test_fn():
return test_tuple[ts]
Comment thread
brandonwillard marked this conversation as resolved.
Resolved

self.assertEqual(test_fn(), test_fn.py_func())

def test_literal_slice_distinct(self):
sl1 = types.misc.SliceLiteral(slice(1, None, None))
sl2 = types.misc.SliceLiteral(slice(None, None, None))
sl3 = types.misc.SliceLiteral(slice(1, None, None))

self.assertNotEqual(sl1, sl2)
self.assertEqual(sl1, sl3)


if __name__ == '__main__':
unittest.main()