inducer / loopy

A code generator for array-based code on CPUs and GPUs
http://mathema.tician.de/software/loopy
MIT License
580 stars 70 forks source link

Math callables not registered inside pymbolic LogicalNot #789

Closed connorjward closed 1 year ago

connorjward commented 1 year ago

The code:

import loopy as lp
import pymbolic as pym

insn = lp.Assignment(
    pym.parse("x[0]"),
    pym.parse("abs(y[0]) > 1").not_(),  # this is the problem line
    within_inames=frozenset({"i"})
)
knl = lp.make_kernel(
    "{ [i]: 0 <= i < 3 }",
    [insn],
    [
        lp.GlobalArg("x", shape=None, dtype=float),
        lp.GlobalArg("y", shape=None, dtype=float),
    ],
    target=lp.CTarget(),
)

device_code = lp.generate_code_v2(knl).device_code()

generates the following error message:

Traceback (most recent call last):
  File "/home/connor/.local/opt/firedrake39/src/firedrake/mytest.py", line 20, in <module>
    device_code = lp.generate_code_v2(knl).device_code()
  File "/home/connor/.local/opt/firedrake39/src/loopy/loopy/codegen/__init__.py", line 598, in genera
te_code_v2
    program = preprocess_program(program)
  File "/home/connor/.local/opt/firedrake39/src/loopy/loopy/tools.py", line 949, in wrapper
    result = func(*args, **kwargs)
  File "/home/connor/.local/opt/firedrake39/src/loopy/loopy/preprocess.py", line 811, in preprocess_p
rogram
    program = infer_unknown_types(program, expect_completion=False)
  File "/home/connor/.local/opt/firedrake39/src/loopy/loopy/type_inference.py", line 1067, in infer_u
nknown_types
    return clbl_inf_ctx.finish_program(program)
  File "/home/connor/.local/opt/firedrake39/src/loopy/loopy/translation_unit.py", line 622, in finish
_program
    old_func_id = history[new_id]
KeyError: 'abs'

The same error is generated for more reasonable uses of abs like pym.primitives.If(pym.parse("abs(y[0]) > 1").not_(), 1, 0) but the example given is the most minimal I could find.

In both cases removing the .not_() causes the error to go away.

inducer commented 1 year ago

Thanks for the report. See #790.

connorjward commented 1 year ago

Great. Thanks for the quick fix!