Skip to content

[OptimizeInstructions] Generalize sign-ext folding #9296

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

Merged
tlively merged 6 commits into from
Oct 28, 2026
Merged

[OptimizeInstructions] Generalize sign-ext folding #9296

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
298000e
[OptimizeInstructions] Generalize sign-ext folding
tlively Oct 21, 2026
87176a6
leftShifts
tlively Oct 23, 2026
4c89c01
fixes
tlively Oct 23, 2026
307b8fe
refactor to helper
tlively Oct 23, 2026
10bb05b
test comment
tlively Oct 26, 2026
9d38749
Merge remote-tracking branch 'origin/main' into opt-inst-1-2-sign-ext…
tlively Oct 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 37 additions & 21 deletions src/ir/properties.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,62 +151,78 @@ inline Literals getLiterals(const Expression* curr) {
// Check if an expression is a sign-extend, and if so, returns the value
// that is extended, otherwise nullptr
inline Expression* getSignExtValue(Expression* curr) {
// We only care about i32s here, and ignore i64s, unreachables, etc.
if (curr->type != Type::i32) {
// We only care about integers here, and ignore unreachables, etc.
if (!curr->type.isInteger()) {
return nullptr;
}
if (auto* unary = curr->dynCast<Unary>()) {
if (unary->op == ExtendS8Int32 || unary->op == ExtendS16Int32) {
return unary->value;
switch (unary->op) {
case ExtendS8Int32:
case ExtendS16Int32:
case ExtendS8Int64:
case ExtendS16Int64:
case ExtendS32Int64:
return unary->value;
default:
return nullptr;
}
return nullptr;
}
using namespace Abstract;
using namespace Match;
int32_t leftShift = 0, rightShift = 0;
int64_t leftShift = 0, rightShift = 0;
Expression* extended = nullptr;
if (matches(curr,
binary(ShrSInt32,
binary(ShlInt32, any(&extended), i32(&leftShift)),
i32(&rightShift))) &&
leftShift == rightShift && leftShift != 0) {
return extended;
binary(ShrS,
binary(Shl, any(&extended), ival(&leftShift)),
ival(&rightShift)))) {
auto leftShifts = Bits::getEffectiveShifts(leftShift, curr->type);
auto rightShifts = Bits::getEffectiveShifts(rightShift, curr->type);
if (leftShifts != 0 && leftShifts == rightShifts) {
return extended;
}
}
Comment thread
tlively marked this conversation as resolved.
Resolved
return nullptr;
}

// gets the size of the sign-extended value
inline Index getSignExtBits(Expression* curr) {
assert(curr->type == Type::i32);
assert(curr->type.isInteger());
if (auto* unary = curr->dynCast<Unary>()) {
switch (unary->op) {
case ExtendS8Int32:
case ExtendS8Int64:
return 8;
case ExtendS16Int32:
case ExtendS16Int64:
return 16;
case ExtendS32Int64:
return 32;
default:
WASM_UNREACHABLE("invalid unary operation");
}
} else {
auto* rightShift = curr->cast<Binary>()->right;
return 32 - Bits::getEffectiveShifts(rightShift);
return curr->type.getByteSize() * 8 - Bits::getEffectiveShifts(rightShift);
}
}

// Check if an expression is almost a sign-extend: perhaps the inner shift
// is too large. We can split the shifts in that case, which is sometimes
// useful (e.g. if we can remove the signext)
inline Expression* getAlmostSignExt(Expression* curr) {
using namespace Abstract;
using namespace Match;
int32_t leftShift = 0, rightShift = 0;
int64_t leftShift = 0, rightShift = 0;
Expression* extended = nullptr;
if (matches(curr,
binary(ShrSInt32,
binary(ShlInt32, any(&extended), i32(&leftShift)),
i32(&rightShift))) &&
Bits::getEffectiveShifts(rightShift, Type::i32) <=
Bits::getEffectiveShifts(leftShift, Type::i32) &&
rightShift != 0) {
return extended;
binary(ShrS,
binary(Shl, any(&extended), ival(&leftShift)),
ival(&rightShift)))) {
auto leftShifts = Bits::getEffectiveShifts(leftShift, curr->type);
auto rightShifts = Bits::getEffectiveShifts(rightShift, curr->type);
if (rightShifts != 0 && rightShifts <= leftShifts) {
return extended;
}
}
return nullptr;
}
Expand Down
90 changes: 48 additions & 42 deletions src/passes/OptimizeInstructions.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -604,21 +604,8 @@ struct OptimizeInstructions
if (auto* ext = Properties::getAlmostSignExt(curr)) {
Index extraLeftShifts;
auto bits = Properties::getAlmostSignExtBits(curr, extraLeftShifts);
if (extraLeftShifts == 0) {
if (auto* load = getFallthrough(ext)->dynCast<Load>()) {
// pattern match a load of 8 bits and a sign extend using a shl of
// 24 then shr_s of 24 as well, etc.
if (LoadUtils::canBeSigned(load) &&
((load->bytes == 1 && bits == 8) ||
(load->bytes == 2 && bits == 16))) {
// if the value falls through, we can't alter the load, as it
// might be captured in a tee
if (load->signed_ == true || load == ext) {
load->signed_ = true;
return replaceCurrent(ext);
}
}
}
if (extraLeftShifts == 0 && trySignExtLoad(ext, bits)) {
return replaceCurrent(ext);
}
// We can in some cases remove part of a sign extend, that is,
// (x << A) >> B => x << (A - B)
Expand All @@ -633,7 +620,8 @@ struct OptimizeInstructions
if (isSignExted(ext, bits) && extraLeftShifts == 0) {
return replaceCurrent(removeAlmostSignExt(curr));
}
} else if (curr->op == EqInt32 || curr->op == NeInt32) {
} else if (curr->op == EqInt32 || curr->op == NeInt32 ||
curr->op == EqInt64 || curr->op == NeInt64) {
if (auto* c = curr->right->dynCast<Const>()) {
if (auto* ext = Properties::getSignExtValue(curr->left)) {
// We are comparing a sign extend to a constant, which means we can
Expand All @@ -649,10 +637,10 @@ struct OptimizeInstructions
// the upper bits, which means that the relevant bits on the left
// side are either all 0, or all 1.
auto bits = Properties::getSignExtBits(curr->left);
uint32_t right = c->value.geti32();
uint32_t numRelevantBits = 32 - bits + 1;
uint32_t setRelevantBits =
Bits::popCount(right >> uint32_t(bits - 1));
uint64_t right = c->value.getUnsigned();
Index totalBits = c->type.getByteSize() * 8;
Index numRelevantBits = totalBits - bits + 1;
Index setRelevantBits = Bits::popCount(right >> (bits - 1));
// If all the relevant bits on C are zero
// then we can mask off the high bits instead of sign-extending x.
// This is valid because if x is negative, then the comparison was
Expand All @@ -670,14 +658,17 @@ struct OptimizeInstructions
// identical, then the upper bits don't matter, so masking them
// off both sides is fine.
curr->left = makeZeroExt(ext, bits);
c->value = c->value.and_(Literal(Bits::lowBitMask(bits)));
c->value = c->value.and_(
Literal::makeFromInt64(Bits::lowBitMask64(bits), c->type));
return replaceCurrent(curr);
} else {
// Otherwise, C's relevant bits are mixed, and then the two sides
// can never be equal, as the left side's bits cannot be mixed.
Builder builder(*getModule());
// The result is either always true, or always false.
c->value = Literal::makeFromInt32(curr->op == NeInt32, c->type);
c->value = Literal::makeFromInt32(
curr->op == NeInt32 || curr->op == NeInt64, Type::i32);
c->type = Type::i32;
Comment thread
kripken marked this conversation as resolved.
Resolved
return replaceCurrent(
builder.makeSequence(builder.makeDrop(ext), c));
}
Expand Down Expand Up @@ -1058,19 +1049,19 @@ struct OptimizeInstructions
}
}

// Simple sign extends can be removed if the value is already sign-extended.
auto signExtBits = getSignExtBits(curr->value);
if (signExtBits > 0) {
// Note that we can handle the case of |curr| having a larger sign-extend:
// if we have an 8-bit value in 32-bit, then there are 24 sign bits, and
// doing a sign-extend to 16 will only affect 16 of those 24, and the
// effect is to leave them as they are.
if ((curr->op == ExtendS8Int32 && signExtBits <= 8) ||
(curr->op == ExtendS16Int32 && signExtBits <= 16) ||
(curr->op == ExtendS8Int64 && signExtBits <= 8) ||
(curr->op == ExtendS16Int64 && signExtBits <= 16) ||
(curr->op == ExtendS32Int64 && signExtBits <= 32)) {
return replaceCurrent(curr->value);
if (auto* ext = Properties::getSignExtValue(curr)) {
auto bits = Properties::getSignExtBits(curr);
if (trySignExtLoad(ext, bits)) {

@kripken kripken Oct 23, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Can we not use getAlmostSignExtBits like the other call to this helper? If not, why must we be more precise here?

@tlively tlively Oct 24, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Two things:

  1. The other call is guarded by extraLeftShifts == 0, so it is also looking at an exact sign extension. The only reason there is a Properties::getAlmostSignExtBits call is because if there are extra left shifts, it falls through and tries other optimizations. The trySignExtLoad optimization only works for exact sign extensions because it is folding the sign extension into a sign-extending load.
  2. getAlmostSignExt and getAlmostSignExtendBits only work with shift-based sign extensions right now. We could generalize them, but there wouldn't be any benefit in this PR due to the previous point.

return replaceCurrent(ext);
}
// Simple sign extends can be removed if the value is already
// sign-extended. Note that we can handle the case of |curr| having a
// larger sign-extend: if we have an 8-bit value in 32-bit, then there are
// 24 sign bits, and doing a sign-extend to 16 will only affect 16 of
// those 24, and the effect is to leave them as they are.
auto signExtBits = getSignExtBits(ext);
if (signExtBits > 0 && signExtBits <= bits) {
return replaceCurrent(ext);
}
Comment on lines +1062 to 1065

@gemini-code-assist gemini-code-assist Bot Oct 23, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

critical

The call to getSignExtBits(ext) is unsafe because ext (which is curr->value) is not guaranteed to be a sign-extend expression. For example, if curr is (i32.extend8_s (local.get $x)), ext will be (local.get $x). Since LocalGet is neither a Unary nor a Binary sign-extend, getSignExtBits will fall through to the else branch and attempt to perform curr->cast<Binary>() on the LocalGet, resulting in an assertion failure/crash.

To fix this, you should check if ext is actually a sign-extend expression before calling getSignExtBits(ext).

      if (Properties::getSignExtValue(ext)) {
        auto signExtBits = getSignExtBits(ext);
        if (signExtBits > 0 && signExtBits <= bits) {
          return replaceCurrent(ext);
        }
      }

}

Expand Down Expand Up @@ -4415,8 +4406,10 @@ struct OptimizeInstructions

Expression* makeZeroExt(Expression* curr, int32_t bits) {
Builder builder(*getModule());
return builder.makeBinary(
AndInt32, curr, builder.makeConst(Literal(Bits::lowBitMask(bits))));
return builder.makeBinary(Abstract::getBinary(curr->type, Abstract::And),
curr,
builder.makeConst(Literal::makeFromInt64(
Bits::lowBitMask64(bits), curr->type)));
}

// given an "almost" sign extend - either a proper one, or it
Expand All @@ -4436,6 +4429,24 @@ struct OptimizeInstructions
return inner;
}

// Given `ext`, the value being sign-extended by a parent sign-extension from
// `bits` low bits (e.g. 8 for `extend8_s` or `(x << 24) >>_s 24`), fold the
// sign-extension into a narrow load if possible.
bool trySignExtLoad(Expression* ext, Index bits) {
if (auto* load = getFallthrough(ext)->dynCast<Load>()) {
// Pattern match a load of 8 bits and a sign extend to 8 bits, etc.
if (LoadUtils::canBeSigned(load) && load->bytes * 8 == bits) {
// If the value falls through, we can't alter the load, as it might be
// captured in a tee.
if (load->signed_ || load == ext) {
load->signed_ = true;
return true;
}
}
}
return false;
}

// Check if an expression is already sign-extended to an exact number of bits.
bool isSignExted(Expression* curr, Index bits) {
return getSignExtBits(curr) == bits;
Expand Down Expand Up @@ -5398,11 +5409,6 @@ struct OptimizeInstructions
// unaryOp(unaryOp(x)) ==> unaryOp(x)
return unaryInner;
}
case ExtendS8Int32:
case ExtendS16Int32: {
assert(getModule()->features.hasSignExt());
return unaryInner;
}
case EqZInt32: {
// eqz(eqz(bool(x))) ==> bool(x)
if (Bits::getMaxBits(unaryInner->value, this) == 1) {
Expand Down
35 changes: 35 additions & 0 deletions test/lit/passes/optimize-instructions-atomics.wast
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,22 @@
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (i32.atomic.load8_u
;; CHECK-NEXT: (i32.const 100)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (i32.extend8_s
;; CHECK-NEXT: (i32.atomic.load8_u
;; CHECK-NEXT: (i32.const 100)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $x
;; An 8-bit shift sign-extend cannot be folded into an atomic load, which
;; cannot be signed, so it is only canonicalized to i32.extend8_s.
(drop
(i32.shr_s
(i32.shl
Expand All @@ -26,6 +40,27 @@
(i32.const 24)
)
)
;; A 12-bit shift sign-extend on an 8-bit unsigned atomic load is a no-op
;; because the upper bits of the 12-bit value are already zero.
(drop
(i32.shr_s
(i32.shl
(i32.atomic.load8_u ;; can't be signed
(i32.const 100)
)
(i32.const 20)
)
(i32.const 20)
)
)
;; A unary sign-extend also cannot be folded into an atomic load.
(drop
(i32.extend8_s
(i32.atomic.load8_u ;; can't be signed
(i32.const 100)
)
)
)
)

;; CHECK: (func $dont_simplify_reinterpret_atomic_load_store (param $x i32) (param $y f32)
Expand Down
Loading
Loading