[llvm] [CodeGenPrepare] Reuse free wider extensions without truncating (PR #228750)

Sat Oct 3 10:42:54 PDT 2026

llvmorg-github-actions[bot] wrote:




@llvm/pr-subscribers-llvm-transforms

Author: Dávid Bolvanský (davidbolvansky)

<details>
<summary>Changes</summary>

Reuse an existing sign/zero extension for a wider extension of the same kind when the target considers it free. This avoids a truncate/re-extend pair that hides an extending-load promotion opportunity.

Motivated by zstd's `ZSTD_decompressSequences`: byte-sized FSE bit counts feed both 32-bit calculations and 64-bit `BIT_mask` indices. With #228736 applied, this improves zstd decompression by **10.1% on a source corpus and 12.6% on JSON** on Apple M1 Pro, versus #228736 alone, in a decoder-only compiler comparison.

Assisted-by: OpenAI Codex


---
Full diff: https://github.com/llvm/llvm-project/pull/228750.diff


2 Files Affected:

- (modified) llvm/lib/CodeGen/CodeGenPrepare.cpp (+16) 
- (added) llvm/test/Transforms/CodeGenPrepare/reuse-extension.ll (+228) 


``````````diff

diff --git a/llvm/lib/CodeGen/CodeGenPrepare.cpp b/llvm/lib/CodeGen/CodeGenPrepare.cpp
index 0b96c77ff33ea..fbb4d3b4e0261 100644
--- a/llvm/lib/CodeGen/CodeGenPrepare.cpp
+++ b/llvm/lib/CodeGen/CodeGenPrepare.cpp
@@ -7391,6 +7391,22 @@ bool CodeGenPrepare::optimizeExtUses(Instruction *I) {
       continue;
 
     // Both src and def are live in this block. Rewrite the use.
+    // Reuse an extension directly when a wider extension of the same kind is
+    // free for its users.
+    // Going through a truncation instead can hide an extending-load
+    // opportunity.
+    if (I->getOpcode() == User->getOpcode() &&
+        User->getType()->getScalarSizeInBits() >
+            I->getType()->getScalarSizeInBits()) {
+      U = I;
+      if (TLI->isExtFree(User)) {
+        ++NumExtUses;
+        MadeChange = true;
+        continue;
+      }
+      U = Src;
+    }
+
     Instruction *&InsertedTrunc = InsertedTruncs[UserBB];
 
     if (!InsertedTrunc) {
diff --git a/llvm/test/Transforms/CodeGenPrepare/reuse-extension.ll b/llvm/test/Transforms/CodeGenPrepare/reuse-extension.ll
new file mode 100644
index 0000000000000..8484ea8b42f7a
--- /dev/null
+++ b/llvm/test/Transforms/CodeGenPrepare/reuse-extension.ll
@@ -0,0 +1,228 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -passes='require<profile-summary>,function(codegenprepare)' -mtriple=aarch64 < %s | FileCheck %s --check-prefixes=CHECK,AARCH64
+; RUN: opt -S -passes='require<profile-summary>,function(codegenprepare)' -mtriple=x86_64 < %s | FileCheck %s --check-prefixes=CHECK,X86
+; REQUIRES: aarch64-registered-target, x86-registered-target
+
+; Reusing a free wider extension must not introduce a truncation that
+; prevents CodeGenPrepare from moving the extension to the original load.
+declare void @consume(i32)
+declare void @consume64(i64)
+
+define i64 @zext_i8(ptr %p, ptr %table, i1 %cold) {
+;
+; CHECK-LABEL: @zext_i8(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[V:%.*]] = load i8, ptr [[P:%.*]], align 1
+; CHECK-NEXT:    [[PROMOTED:%.*]] = zext i8 [[V]] to i64
+; CHECK-NEXT:    [[V32:%.*]] = zext i8 [[V]] to i32
+; CHECK-NEXT:    br i1 [[COLD:%.*]], label [[CALL:%.*]], label [[JOIN:%.*]]
+; CHECK:       call:
+; CHECK-NEXT:    call void @consume(i32 [[V32]])
+; CHECK-NEXT:    br label [[JOIN]]
+; CHECK:       join:
+; CHECK-NEXT:    [[PTR:%.*]] = getelementptr i32, ptr [[TABLE:%.*]], i64 [[PROMOTED]]
+; CHECK-NEXT:    [[MASK:%.*]] = load i32, ptr [[PTR]], align 4
+; CHECK-NEXT:    [[SUM:%.*]] = add i32 [[MASK]], [[V32]]
+; CHECK-NEXT:    [[RESULT:%.*]] = zext i32 [[SUM]] to i64
+; CHECK-NEXT:    ret i64 [[RESULT]]
+;
+entry:
+  %v = load i8, ptr %p
+  %v32 = zext i8 %v to i32
+  br i1 %cold, label %call, label %join
+call:
+  call void @consume(i32 %v32)
+  br label %join
+join:
+  %v64 = zext i8 %v to i64
+  %ptr = getelementptr i32, ptr %table, i64 %v64
+  %mask = load i32, ptr %ptr
+  %sum = add i32 %mask, %v32
+  %result = zext i32 %sum to i64
+  ret i64 %result
+}
+
+define i64 @zext_i16(ptr %p, ptr %table, i1 %cold) {
+;
+; CHECK-LABEL: @zext_i16(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[V:%.*]] = load i16, ptr [[P:%.*]], align 2
+; CHECK-NEXT:    [[PROMOTED:%.*]] = zext i16 [[V]] to i64
+; CHECK-NEXT:    [[V32:%.*]] = zext i16 [[V]] to i32
+; CHECK-NEXT:    br i1 [[COLD:%.*]], label [[CALL:%.*]], label [[JOIN:%.*]]
+; CHECK:       call:
+; CHECK-NEXT:    call void @consume(i32 [[V32]])
+; CHECK-NEXT:    br label [[JOIN]]
+; CHECK:       join:
+; CHECK-NEXT:    [[PTR:%.*]] = getelementptr i32, ptr [[TABLE:%.*]], i64 [[PROMOTED]]
+; CHECK-NEXT:    [[MASK:%.*]] = load i32, ptr [[PTR]], align 4
+; CHECK-NEXT:    [[SUM:%.*]] = add i32 [[MASK]], [[V32]]
+; CHECK-NEXT:    [[RESULT:%.*]] = zext i32 [[SUM]] to i64
+; CHECK-NEXT:    ret i64 [[RESULT]]
+;
+entry:
+  %v = load i16, ptr %p
+  %v32 = zext i16 %v to i32
+  br i1 %cold, label %call, label %join
+call:
+  call void @consume(i32 %v32)
+  br label %join
+join:
+  %v64 = zext i16 %v to i64
+  %ptr = getelementptr i32, ptr %table, i64 %v64
+  %mask = load i32, ptr %ptr
+  %sum = add i32 %mask, %v32
+  %result = zext i32 %sum to i64
+  ret i64 %result
+}
+
+define i64 @sext_then_zext(ptr %p, ptr %table, i1 %cold) {
+;
+; CHECK-LABEL: @sext_then_zext(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[V:%.*]] = load i8, ptr [[P:%.*]], align 1
+; CHECK-NEXT:    [[V32:%.*]] = sext i8 [[V]] to i32
+; CHECK-NEXT:    br i1 [[COLD:%.*]], label [[CALL:%.*]], label [[JOIN:%.*]]
+; CHECK:       call:
+; CHECK-NEXT:    call void @consume(i32 [[V32]])
+; CHECK-NEXT:    br label [[JOIN]]
+; CHECK:       join:
+; CHECK-NEXT:    [[TMP0:%.*]] = trunc i32 [[V32]] to i8
+; CHECK-NEXT:    [[V64:%.*]] = zext i8 [[TMP0]] to i64
+; CHECK-NEXT:    [[PTR:%.*]] = getelementptr i32, ptr [[TABLE:%.*]], i64 [[V64]]
+; CHECK-NEXT:    [[MASK:%.*]] = load i32, ptr [[PTR]], align 4
+; CHECK-NEXT:    [[SUM:%.*]] = add i32 [[MASK]], [[V32]]
+; CHECK-NEXT:    [[RESULT:%.*]] = zext i32 [[SUM]] to i64
+; CHECK-NEXT:    ret i64 [[RESULT]]
+;
+entry:
+  %v = load i8, ptr %p
+  %v32 = sext i8 %v to i32
+  br i1 %cold, label %call, label %join
+call:
+  call void @consume(i32 %v32)
+  br label %join
+join:
+  %v64 = zext i8 %v to i64
+  %ptr = getelementptr i32, ptr %table, i64 %v64
+  %mask = load i32, ptr %ptr
+  %sum = add i32 %mask, %v32
+  %result = zext i32 %sum to i64
+  ret i64 %result
+}
+
+define i64 @zext_then_sext(ptr %p, ptr %table, i1 %cold) {
+;
+; CHECK-LABEL: @zext_then_sext(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[V:%.*]] = load i8, ptr [[P:%.*]], align 1
+; CHECK-NEXT:    [[V32:%.*]] = zext i8 [[V]] to i32
+; CHECK-NEXT:    br i1 [[COLD:%.*]], label [[CALL:%.*]], label [[JOIN:%.*]]
+; CHECK:       call:
+; CHECK-NEXT:    call void @consume(i32 [[V32]])
+; CHECK-NEXT:    br label [[JOIN]]
+; CHECK:       join:
+; CHECK-NEXT:    [[TMP0:%.*]] = trunc i32 [[V32]] to i8
+; CHECK-NEXT:    [[V64:%.*]] = sext i8 [[TMP0]] to i64
+; CHECK-NEXT:    [[PTR:%.*]] = getelementptr i32, ptr [[TABLE:%.*]], i64 [[V64]]
+; CHECK-NEXT:    [[MASK:%.*]] = load i32, ptr [[PTR]], align 4
+; CHECK-NEXT:    [[SUM:%.*]] = add i32 [[MASK]], [[V32]]
+; CHECK-NEXT:    [[RESULT:%.*]] = zext i32 [[SUM]] to i64
+; CHECK-NEXT:    ret i64 [[RESULT]]
+;
+entry:
+  %v = load i8, ptr %p
+  %v32 = zext i8 %v to i32
+  br i1 %cold, label %call, label %join
+call:
+  call void @consume(i32 %v32)
+  br label %join
+join:
+  %v64 = sext i8 %v to i64
+  %ptr = getelementptr i32, ptr %table, i64 %v64
+  %mask = load i32, ptr %ptr
+  %sum = add i32 %mask, %v32
+  %result = zext i32 %sum to i64
+  ret i64 %result
+}
+
+
+define i64 @sext_address(ptr %p, ptr %table, i1 %cold) {
+; AARCH64-LABEL: @sext_address(
+; AARCH64-NEXT:  entry:
+; AARCH64-NEXT:    [[V:%.*]] = load i8, ptr [[P:%.*]], align 1
+; AARCH64-NEXT:    [[V64:%.*]] = sext i8 [[V]] to i64
+; AARCH64-NEXT:    [[V32:%.*]] = sext i8 [[V]] to i32
+; AARCH64-NEXT:    br i1 [[COLD:%.*]], label [[CALL:%.*]], label [[JOIN:%.*]]
+; AARCH64:       call:
+; AARCH64-NEXT:    call void @consume(i32 [[V32]])
+; AARCH64-NEXT:    br label [[JOIN]]
+; AARCH64:       join:
+; AARCH64-NEXT:    [[PTR:%.*]] = getelementptr i32, ptr [[TABLE:%.*]], i64 [[V64]]
+; AARCH64-NEXT:    [[MASK:%.*]] = load i32, ptr [[PTR]], align 4
+; AARCH64-NEXT:    [[SUM:%.*]] = add i32 [[MASK]], [[V32]]
+; AARCH64-NEXT:    [[RESULT:%.*]] = zext i32 [[SUM]] to i64
+; AARCH64-NEXT:    ret i64 [[RESULT]]
+;
+; X86-LABEL: @sext_address(
+; X86-NEXT:  entry:
+; X86-NEXT:    [[V:%.*]] = load i8, ptr [[P:%.*]], align 1
+; X86-NEXT:    [[V32:%.*]] = sext i8 [[V]] to i32
+; X86-NEXT:    br i1 [[COLD:%.*]], label [[CALL:%.*]], label [[JOIN:%.*]]
+; X86:       call:
+; X86-NEXT:    call void @consume(i32 [[V32]])
+; X86-NEXT:    br label [[JOIN]]
+; X86:       join:
+; X86-NEXT:    [[TMP0:%.*]] = trunc i32 [[V32]] to i8
+; X86-NEXT:    [[V64:%.*]] = sext i8 [[TMP0]] to i64
+; X86-NEXT:    [[PTR:%.*]] = getelementptr i32, ptr [[TABLE:%.*]], i64 [[V64]]
+; X86-NEXT:    [[MASK:%.*]] = load i32, ptr [[PTR]], align 4
+; X86-NEXT:    [[SUM:%.*]] = add i32 [[MASK]], [[V32]]
+; X86-NEXT:    [[RESULT:%.*]] = zext i32 [[SUM]] to i64
+; X86-NEXT:    ret i64 [[RESULT]]
+;
+entry:
+  %v = load i8, ptr %p
+  %v32 = sext i8 %v to i32
+  br i1 %cold, label %call, label %join
+call:
+  call void @consume(i32 %v32)
+  br label %join
+join:
+  %v64 = sext i8 %v to i64
+  %ptr = getelementptr i32, ptr %table, i64 %v64
+  %mask = load i32, ptr %ptr
+  %sum = add i32 %mask, %v32
+  %result = zext i32 %sum to i64
+  ret i64 %result
+}
+
+define i64 @sext_standalone(ptr %p, ptr %table, i1 %cold) {
+; CHECK-LABEL: @sext_standalone(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[V:%.*]] = load i8, ptr [[P:%.*]], align 1
+; CHECK-NEXT:    [[V32:%.*]] = sext i8 [[V]] to i32
+; CHECK-NEXT:    br i1 [[COLD:%.*]], label [[CALL:%.*]], label [[JOIN:%.*]]
+; CHECK:       call:
+; CHECK-NEXT:    call void @consume(i32 [[V32]])
+; CHECK-NEXT:    br label [[JOIN]]
+; CHECK:       join:
+; CHECK-NEXT:    [[TMP0:%.*]] = trunc i32 [[V32]] to i8
+; CHECK-NEXT:    [[V64:%.*]] = sext i8 [[TMP0]] to i64
+; CHECK-NEXT:    call void @consume64(i64 [[V64]])
+; CHECK-NEXT:    [[RESULT:%.*]] = zext i32 [[V32]] to i64
+; CHECK-NEXT:    ret i64 [[RESULT]]
+;
+entry:
+  %v = load i8, ptr %p
+  %v32 = sext i8 %v to i32
+  br i1 %cold, label %call, label %join
+call:
+  call void @consume(i32 %v32)
+  br label %join
+join:
+  %v64 = sext i8 %v to i64
+  call void @consume64(i64 %v64)
+  %result = zext i32 %v32 to i64
+  ret i64 %result
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/228750
