diff --git a/loopy/target/execution.py b/loopy/target/execution.py index 2aa76e099d8e50a2949c616736b30f725fb10bb4..facd56a071d3e01ec49a92d4a5b937d99d1eb5a1 100644 --- a/loopy/target/execution.py +++ b/loopy/target/execution.py @@ -363,6 +363,10 @@ class ExecutionWrapperGeneratorBase(object): from loopy.types import NumpyType gen("# {{{ set up array arguments") + + gen("") + gen("def _lpy_filter_stride(shape, stride):") + gen(" return tuple(s for dim, s in zip(shape, stride) if dim > 1)") gen("") if not options.no_numpy: @@ -516,13 +520,21 @@ class ExecutionWrapperGeneratorBase(object): itemsize = kernel_arg.dtype.numpy_dtype.itemsize sym_strides = tuple( itemsize*s_i for s_i in arg.unvec_strides) - gen("if %s.strides != %s:" - % (arg.name, strify(sym_strides))) + gen("if _lpy_filter_stride(%s.shape, %s.strides) != " + "_lpy_filter_stride(%s.shape, %s):" + % ( + arg.name, arg.name, arg.name, + strify(sym_strides))) with Indentation(gen): gen("raise TypeError(\"strides mismatch on " - "argument '%s' (got: %%s, expected: %%s)\" " - "%% (%s.strides, %s))" - % (arg.name, arg.name, strify(sym_strides))) + "argument '%s' " + "(after removing unit length dims, " + "got: %%s, expected: %%s)\" " + "%% (_lpy_filter_stride(%s.shape, %s.strides), " + "_lpy_filter_stride(%s.shape, %s)))" + % ( + arg.name, arg.name, arg.name, arg.name, + strify(sym_strides))) if not arg.allows_offset: gen("if hasattr(%s, 'offset') and %s.offset:" % ( diff --git a/test/test_loopy.py b/test/test_loopy.py index e624ed346cd696bf18a116e9373f8e765dafdc9a..375b59dcb38bb889124baadd6cc473fe029d48c6 100644 --- a/test/test_loopy.py +++ b/test/test_loopy.py @@ -2746,6 +2746,26 @@ def test_arg_inference_for_predicates(): assert knl.arg_dict["incr"].shape == (10,) +def test_relaxed_stride_checks(ctx_factory): + # Check that loopy is compatible with numpy's relaxed stride rules. + ctx = ctx_factory() + + knl = lp.make_kernel("{[i,j]: 0 <= i <= n and 0 <= j <= m}", + """ + a[i] = sum(j, A[i,j] * b[j]) + """) + + with cl.CommandQueue(ctx) as queue: + A = np.zeros((1, 10), order="F") + # Force convert A to C order. numpy will preserve strides in this case. + A = np.array(A, copy=False, order="C") + b = np.zeros(10, dtype=np.float64) + + evt, (a,) = knl(queue, A=A, b=b) + + assert a == 0 + + def test_add_prefetch_works_in_lhs_index(): knl = lp.make_kernel( "{ [n,k,l,k1,l1,k2,l2]: "