G
YouTube TikTok Games
Skip to content

[CuTeDSL] Keep -0.0 and 0.0 constants apart in the IR constant cache - #3849

Open
Zhonglin1231 wants to merge 1 commit into
NVIDIA:mainfrom
Zhonglin1231:cutedsl-negative-zero-const
Open

Zhonglin1231 wants to merge 1 commit into
NVIDIA:mainfrom
Zhonglin1231:cutedsl-negative-zero-const

Conversation

@Zhonglin1231

Copy link
Copy Markdown

Fixes #3845.

const() goes through _const_cached, an lru_cache keyed on the Python value. 0.0 and -0.0 compare equal and hash alike, so whichever zero a kernel materialized first came back for the other one too. This builds -0.0 through the uncached function (__wrapped__, set by functools.wraps in lru_cache_ir), so it can neither hit nor seed the +0.0 entry. Every other value is cached as before.

math in this module is the MLIR dialect, so the stdlib copysign is imported as _copysign.

Test

test/python/CuTeDSL/test_negative_zero_constant.py stores Float32(0.0) and Float32(-0.0) in one kernel, in both orders, and checks the sign bits. It needs a CUDA GPU and skips without one.

On an A10G with nvidia-cutlass-dsl 4.8.1, both orders fail without the fix (sign bits [False, False] and [True, True]) and pass with it.

const() goes through an lru_cache keyed on the Python value. 0.0 and
-0.0 compare equal and hash alike, so whichever zero a kernel
materialized first came back for the other one too: storing
Float32(0.0) and then Float32(-0.0) wrote +0.0 twice. Build -0.0
uncached so it can neither hit nor seed the +0.0 entry.

Co-authored-by: Ruilin Gao <229880965+ruiling-smartbear@users.noreply.github.com>
Signed-off-by: Zhong Lin <125393049+Zhonglin1231@users.noreply.github.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] CuTe DSL: -0.0 and 0.0 constants in one kernel come out as the same zero

2 participants