Skip to content

[stdlib] Fix SIMD[..., u128/s128] construction from UInt #4339

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[stdlib] Fix SIMD[..., u128/s128] construction from UInt #4339

wants to merge 1 commit into from
+17 −1

Conversation

soraros
Copy link
Contributor

@soraros soraros commented Apr 12, 2025

Close #4338.

@soraros soraros requested a review from a team as a code owner April 12, 2025 04:34
soraros
soraros commented Apr 12, 2025
View reviewed changes
mojo/stdlib/test/builtin/test_simd.mojo Outdated
Comment on lines 109 to 112
alias a = rebind[UInt](-1)
alias a_str = String(a)
assert_equal(a_str, String(UInt128(a)))
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can't do much better than stringify due to #4049.

soraros
soraros commented Apr 12, 2025
View reviewed changes
mojo/stdlib/src/builtin/simd.mojo Outdated
self = Self(value.value)
@parameter
if bitwidthof[dtype]() > bitwidthof[DType.index]():
self = Self(rebind[UInt64](value))
Copy link
Contributor Author

@soraros soraros Apr 12, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cannot use UInt64(value) here. I think there's a compiler bug(?) where v.cast[index]().cast[int128]() gets simplified to v.cast[int128]() (on the pop level), which lowers to sext yet again. So we use rebind to skip the constructor from index entirely.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actually, it's probably not really a compiler bug. Suppose the target type is UInt128. If we go through UInt64, we get roughly the following MLIR:

var v: index = ...
var a = pop.cast[int64](v)
var b = pop.cast[int128](a)

Since index is treated as signed (?), the first cast is a no-op, and the second is a sign extension.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Question: Aren't you assuming 64 bitwidth machines by using UInt64?

Copy link
Contributor Author

@soraros soraros Apr 13, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yea, I'm being lazy here. Will fix it properly when the team confirmes that rebind is an acceptable approach.

Copy link
Contributor Author

@soraros soraros Apr 15, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actually, rebind[T] doesn't work when alias T = Scalar[_unsigned_integral_type_of[DType.index]()], so I opted for a different fix.

@soraros soraros force-pushed the index-to-i128 branch 2 times, most recently from fb3be68 to e883e8b Compare April 12, 2025 15:21
martinvuyk
martinvuyk reviewed Apr 13, 2025
View reviewed changes
@soraros soraros force-pushed the index-to-i128 branch 3 times, most recently from 4596067 to dca848d Compare April 15, 2025 16:48
@soraros soraros force-pushed the index-to-i128 branch 3 times, most recently from 2ef9723 to 753d223 Compare April 27, 2025 15:55
@soraros soraros force-pushed the index-to-i128 branch 4 times, most recently from 508e0f1 to f61e0e9 Compare May 9, 2025 05:17
@soraros
Copy link
Contributor Author

soraros commented May 16, 2025

@JoeLoser Could you take a look at this one, since it's a bit insidious, and maybe related to Gabriel's ongoing atof work.

@JoeLoser
Copy link
Collaborator

!sync

@JoeLoser
Copy link
Collaborator

@JoeLoser Could you take a look at this one, since it's a bit insidious, and maybe related to Gabriel's ongoing atof work.

LGTM, thanks! Syncing now.

@modularbot modularbot added the imported-internally Signals that a given pull request has been imported internally. label May 16, 2025
@JoeLoser
Copy link
Collaborator

JoeLoser commented May 16, 2025 •
edited
Loading

@soraros FYI seeing failures internally like:

open-source/max/mojo/stdlib/stdlib/memory/unsafe.mojo:82:6: error: 'pop.bitcast' op operand type '!pop.scalar<index>' and result type '!pop.scalar<ui64>' are cast incompatible
open-source/max/mojo/stdlib/stdlib/builtin/simd.mojo:438:36: note: called from
open-source/max/mojo/stdlib/stdlib/memory/unsafe.mojo:82:6: note: see current operation: %4 = "pop.bitcast"(%0) : (!pop.scalar<index>) -> !pop.scalar<ui64>
open-source/max/mojo/stdlib/stdlib/memory/unsafe.mojo:82:6: error: 'pop.bitcast' op operand type '!pop.scalar<index>' and result type '!pop.scalar<ui64>' are cast incompatible
open-source/max/mojo/stdlib/stdlib/builtin/simd.mojo:438:36: note: called from
open-source/max/mojo/stdlib/stdlib/memory/unsafe.mojo:82:6: note: see current operation: %4 = "pop.bitcast"(%0) : (!pop.scalar<index>) -> !pop.scalar<ui64>
open-source/max/mojo/stdlib/stdlib/memory/unsafe.mojo:82:6: error: 'pop.bitcast' op operand type '!pop.scalar<index>' and result type '!pop.scalar<ui64>' are cast incompatible
open-source/max/mojo/stdlib/stdlib/builtin/simd.mojo:438:36: note: called from
open-source/max/mojo/stdlib/stdlib/memory/unsafe.mojo:82:6: note: see current operation: %4 = "pop.bitcast"(%0) : (!pop.scalar<index>) -> !pop.scalar<ui64>
open-source/max/mojo/stdlib/stdlib/memory/unsafe.mojo:82:6: error: 'pop.bitcast' op operand type '!pop.scalar<index>' and result type '!pop.scalar<ui64>' are cast incompatible
open-source/max/mojo/stdlib/stdlib/builtin/simd.mojo:438:36: note: called from
open-source/max/mojo/stdlib/stdlib/memory/unsafe.mojo:82:6: note: see current operation: %4 = "pop.bitcast"(%0) : (!pop.scalar<index>) -> !pop.scalar<ui64>

Can you rebase this PR to see if they show up in OSS build for you as well?

@JoeLoser JoeLoser added the blocked Blocked by another issue that must be resolved first label May 22, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@martinvuyk martinvuyk martinvuyk left review comments

@JoeLoser JoeLoser JoeLoser approved these changes

Assignees

@JoeLoser JoeLoser

Labels
blocked Blocked by another issue that must be resolved first imported-internally Signals that a given pull request has been imported internally.
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

[BUG] Mojo doesn't handle UInt to SIMD with bit widths > 64 correctly
4 participants
@soraros @JoeLoser @martinvuyk @modularbot