diff --git a/loopy/kernel/array.py b/loopy/kernel/array.py index a02fc58d97f370d45f36a465c38fa3caf3da9d41..f3f4e9d7338fb83a98fb636bd97158728eab1563 100644 --- a/loopy/kernel/array.py +++ b/loopy/kernel/array.py @@ -643,8 +643,9 @@ class ArrayBase(ImmutableRecord): :arg offset: Offset from the beginning of the buffer to the point from which the strides are counted. May be one of - * 0 + * 0 or None * a string (that is interpreted as an argument name). + * a pymbolic expression * :class:`loopy.auto`, in which case an offset argument is added automatically, immediately following this argument. :class:`loopy.CompiledKernel` is even smarter in its treatment of @@ -1036,7 +1037,9 @@ class ArrayBase(ImmutableRecord): is_written=is_written) - if self.offset: + import loopy as lp + + if self.offset is lp.auto: offset_name = full_name+"_offset" yield ImplementedDataInfo( target=target, @@ -1202,12 +1205,16 @@ def get_access_info(target, ary, index, eval_expr, vectorization_info): return result def apply_offset(sub): - if ary.offset: - offset_name = ary.offset - if offset_name is lp.auto: - offset_name = array_name+"_offset" + import loopy as lp - return var(offset_name) + sub + if ary.offset: + if ary.offset is lp.auto: + return var(array_name+"_offset") + sub + elif isinstance(ary.offset, str): + return var(ary.offset) + sub + else: + # assume it's an expression + return ary.offset + sub else: return sub