diff --git a/examples/python/sparse.py b/examples/python/sparse.py new file mode 100644 index 0000000000000000000000000000000000000000..5fba194eff1195f73f494a9f32112f2934d8fe3a --- /dev/null +++ b/examples/python/sparse.py @@ -0,0 +1,14 @@ +import loopy as lp + +k = lp.make_kernel([ + "[m] -> { [i] : 0 <= i < m }", + "[length] -> { [j] : 0 <= j < length }"], + """ + rowstart = rowstarts[i] + rowend = rowstarts[1 + i] + length = rowend + (-1)*rowstart + rowsum = 0 {id=zerosum} + rowsum = rowsum + x[-1 + colindices[-1 + rowstart + j]]*values[-1 + rowstart + j] {dep=zerosum} + y[i] = rowsum + """) +print(k) diff --git a/loopy/kernel/creation.py b/loopy/kernel/creation.py index 6a5d523e61c27c068a380b7de6c206147d7bb271..830fe0efadb0425673a9c6ac4f62050d384dcd3f 100644 --- a/loopy/kernel/creation.py +++ b/loopy/kernel/creation.py @@ -914,13 +914,27 @@ def guess_arg_shape_if_requested(kernel, default_order): if armap.access_range is None: if armap.bad_subscripts: - raise RuntimeError("cannot determine access range for '%s': " - "undetermined index in subscript(s) '%s'" - % (arg.name, ", ".join( - str(i) for i in armap.bad_subscripts))) + n_axes_in_subscripts = set( + len(sub.index_tuple) for sub in armap.bad_subscripts) - # no subscripts found, let's call it a scalar - shape = () + if len(n_axes_in_subscripts) != 1: + raise RuntimeError("subscripts of '%s' with differing " + "numbers of axes were found" % arg.name) + + n_axes, = n_axes_in_subscripts + + if n_axes == 1: + # Leave shape undetermined--we can live with that for 1D. + shape = (None,) + else: + raise RuntimeError("cannot determine access range for '%s': " + "undetermined index in subscript(s) '%s'" + % (arg.name, ", ".join( + str(i) for i in armap.bad_subscripts))) + + else: + # no subscripts found, let's call it a scalar + shape = () else: from loopy.isl_helpers import static_max_of_pw_aff from loopy.symbolic import pw_aff_to_expr