diff --git a/examples/python/sparse.py b/examples/python/sparse.py index 14d157b6a6bb7bf4601852d5cf344248a9e82a9b..5fba194eff1195f73f494a9f32112f2934d8fe3a 100644 --- a/examples/python/sparse.py +++ b/examples/python/sparse.py @@ -10,9 +10,5 @@ k = lp.make_kernel([ rowsum = 0 {id=zerosum} rowsum = rowsum + x[-1 + colindices[-1 + rowstart + j]]*values[-1 + rowstart + j] {dep=zerosum} y[i] = rowsum - """, - [ - lp.GlobalArg("colindices,values,x", shape=None), - "..." - ]) + """) 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