loopy icon indicating copy to clipboard operation
loopy copied to clipboard

`buffer_array` fails with an assertion error

Open inducer opened this issue 2 years ago • 0 comments

import loopy as lp

knl = lp.make_kernel(
        "{[i, j]: 0<=i, j<10}",
        """
        <> tmp[j] = 2*x[i, j]
        out[i, j] = 2 * tmp[j-1] {if=j>1}
        """, lang_version=(2018, 2))

knl = lp.prioritize_loops(knl, ("i", "j"))
knl = lp.add_dtypes(knl, {"x": "float32"})
knl = lp.buffer_array(knl, "tmp", "j")
print(lp.generate_code_v2(knl).device_code())

Shared by @kaushikcfd.

Fails with

Traceback (most recent call last):
  File "/home/andreas/tmp/fail_buffer_array.py", line 13, in <module>
    knl = lp.buffer_array(knl, "tmp", "j")
  File "/home/andreas/src/loopy/loopy/tools.py", line 926, in wrapper
    result = func(*args, **kwargs)
  File "/home/andreas/src/loopy/loopy/transform/buffer.py", line 518, in buffer_array
    subkernel=buffer_array_for_single_kernel(clbl.subkernel,
  File "/home/andreas/src/loopy/loopy/transform/buffer.py", line 403, in buffer_array_for_single_kernel
    kernel = rule_mapping_context.finish_kernel(aar.map_kernel(kernel))
  File "/home/andreas/src/loopy/loopy/symbolic.py", line 1389, in map_kernel
    new_insns = [
  File "/home/andreas/src/loopy/loopy/symbolic.py", line 1395, in <listcomp>
    insn.with_transformed_expressions(
  File "/home/andreas/src/loopy/loopy/kernel/instruction.py", line 860, in with_transformed_expressions
    expression = f(self.expression)
  File "/home/andreas/src/loopy/loopy/symbolic.py", line 1396, in <lambda>
    lambda expr: self(expr, kernel, insn)))
  File "/home/andreas/src/loopy/loopy/symbolic.py", line 1377, in __call__
    return super().__call__(expr,
  File "/home/andreas/src/pymbolic/pymbolic/mapper/__init__.py", line 244, in __call__
    result = super().rec(expr, *args, **kwargs)
  File "/home/andreas/src/pymbolic/pymbolic/mapper/__init__.py", line 153, in __call__
    return method(expr, *args, **kwargs)
  File "/home/andreas/src/pymbolic/pymbolic/mapper/__init__.py", line 507, in map_product
    children = [self.rec(child, *args, **kwargs) for child in expr.children]
  File "/home/andreas/src/pymbolic/pymbolic/mapper/__init__.py", line 507, in <listcomp>
    children = [self.rec(child, *args, **kwargs) for child in expr.children]
  File "/home/andreas/src/pymbolic/pymbolic/mapper/__init__.py", line 244, in __call__
    result = super().rec(expr, *args, **kwargs)
  File "/home/andreas/src/pymbolic/pymbolic/mapper/__init__.py", line 153, in __call__
    return method(expr, *args, **kwargs)
  File "/home/andreas/src/loopy/loopy/transform/buffer.py", line 80, in map_subscript
    return super().map_subscript(expr, expn_state)
  File "/home/andreas/src/pymbolic/pymbolic/mapper/__init__.py", line 485, in map_subscript
    aggregate = self.rec(expr.aggregate, *args, **kwargs)
  File "/home/andreas/src/pymbolic/pymbolic/mapper/__init__.py", line 244, in __call__
    result = super().rec(expr, *args, **kwargs)
  File "/home/andreas/src/pymbolic/pymbolic/mapper/__init__.py", line 153, in __call__
    return method(expr, *args, **kwargs)
  File "/home/andreas/src/loopy/loopy/transform/buffer.py", line 63, in map_variable
    result = self.map_array_access((), expn_state)
  File "/home/andreas/src/loopy/loopy/transform/buffer.py", line 97, in map_array_access
    assert len(index) == len(abm.non1_storage_axis_flags)
AssertionError

on 8ca762f8.

inducer avatar Jun 06 '22 16:06 inducer