Enzyme icon indicating copy to clipboard operation
Enzyme copied to clipboard

Zig standard library pow function causes enzyme to fail

Open Copper280z opened this issue 5 months ago • 4 comments

Hi, I've got some basic functionality working with Zig using the C API/ABI, but I've noticed that using Zig's std.math.pow function emits some IR that seems to break Enzyme. std.math.cos works fine, as do basic math operations.

This is on an M2 mac, and Enzyme installed via homebrew. I believe the install brought llvm-19 with it, so I'm running Enzyme with that. The error message is huge, so I'll just post the last bit here and attach the input.ll which produced it.

input.txt

 constantinst[  %common.ret.op.i = phi float [ %29, %26 ], [ %48, %47 ], [ %53, %51 ], [ %65, %62 ], [ 0x7FF0000000000000, %24 ], [ 1.000000e+00, %30 ], [ %.1.i, %35 ], [ 0.000000e+00, %43 ], [ 0x7FF0000000000000, %45 ], [ 0x7FF8000000000000, %54 ], [ %163, %160 ], [ %141, %139 ], [ %149, %145 ], [ %159, %152 ], [ %169, %167 ], [ %191, %172 ], [ %171, %170 ], [ %.sroa.09.3.i, %118 ], [ %.sroa.09.3.i, %123 ], [ 0x7FF0000000000000, %.thread ], [ 0.000000e+00, %math.pow.isOddInteger.exit.thread ], [ %spec.select, %25 ], [ %common.ret.op.i.ph, %math.pow.pow__anon_3229.exit.loopexit ]] = 0 val:0 type: {[-1]:Float@float}
 constantinst[  ret float %common.ret.op.i] = 1 val:1 type: {}
cannot handle unknown binary operator:   %74 = and i32 %68, -2139095041

A full repro can be found here (https://github.com/Copper280z/Enzyme-Zig/blob/ef1b186aa42e3cd33ffbfe217708b26a0fd5a37f/src/main.zig#L67), enzyme is run from the buid system, so you can just zig build and it'll build the source, run enzyme, and compile the output to an executable. Uncomment these lines and the enzyme step fails. Everything works as I'd expect with them commented out.

Copper280z avatar Jul 21 '25 00:07 Copper280z

This is a common issue with low-level math functionality. You will want to either call an llvm intrinsic or define a custom rule for that math operator?

vchuravy avatar Jul 21 '25 09:07 vchuravy

Ah, I didn't realize that. That's easy enough to implement on my end, thank you!

Copper280z avatar Jul 21 '25 09:07 Copper280z

so if it's a language std library, I think that's widely used enough that it makes sense to mark it as having "sin"-like functionality within Enzyme proper, see https://github.com/EnzymeAD/Enzyme/blob/d458ef0ac900fdcd49592238a19a91158de1366f/enzyme/Enzyme/PreserveNVVM.cpp#L779 for example where we have a pass doing so for rocm/nvvm (cuda)

wsmoses avatar Jul 22 '25 19:07 wsmoses

I think that would be a great way to handle it from a users perspective, especially as most of the stdlib math functions are generic over type (https://github.com/ziglang/zig/blob/fc4b7c968afa6fa0780a011f3c8cfeaea38b7b98/lib/std/math/pow.zig#L33). However, given the language's instability (pre-1.0), there could (will) be breaking changes.

I tried annotating with metadata, but because the stdlib function is a generic there's no pointer available, so I don't think that can work. For example here's the name pulled from the IR for std.math.pow(f32, ...):math.pow.pow__anon_19626. I think if there's a workaround for this, the C API can be totally usable.

Copper280z avatar Jul 27 '25 14:07 Copper280z