From eda8e999dd01e9a741a43330b693cff94d4cb955 Mon Sep 17 00:00:00 2001
From: luxufan <luxufan@iscas.ac.cn>
Date: Sun, 8 Jan 2023 14:27:43 +0800
Subject: [PATCH] [InstCombine] Combine (zext a) mul (zext b) to
 llvm.umul.with.overflow only if mul has NUW flag

Fixes: https://github.com/llvm/llvm-project/issues/59836

Reviewed By: nikic

Differential Revision: https://reviews.llvm.org/D141031
---
 .../InstCombine/InstCombineCompares.cpp       |  4 +-
 .../Transforms/InstCombine/overflow-mul.ll    | 38 +++++++++++++++++++
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
index 2d9a27a5eacfa..b0ebfcf919668 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -6338,11 +6338,11 @@ Instruction *InstCombinerImpl::visitICmpInst(ICmpInst &I) {
     }
 
     // (zext a) * (zext b)  --> llvm.umul.with.overflow.
-    if (match(Op0, m_Mul(m_ZExt(m_Value(A)), m_ZExt(m_Value(B))))) {
+    if (match(Op0, m_NUWMul(m_ZExt(m_Value(A)), m_ZExt(m_Value(B))))) {
       if (Instruction *R = processUMulZExtIdiom(I, Op0, Op1, *this))
         return R;
     }
-    if (match(Op1, m_Mul(m_ZExt(m_Value(A)), m_ZExt(m_Value(B))))) {
+    if (match(Op1, m_NUWMul(m_ZExt(m_Value(A)), m_ZExt(m_Value(B))))) {
       if (Instruction *R = processUMulZExtIdiom(I, Op1, Op0, *this))
         return R;
     }
diff --git a/llvm/test/Transforms/InstCombine/overflow-mul.ll b/llvm/test/Transforms/InstCombine/overflow-mul.ll
index f80431c871475..96d754f27d34a 100644
--- a/llvm/test/Transforms/InstCombine/overflow-mul.ll
+++ b/llvm/test/Transforms/InstCombine/overflow-mul.ll
@@ -241,3 +241,41 @@ define i1 @pr21445(i8 %a) {
   %cmp = icmp ne i32 %mul, %and
   ret i1 %cmp
 }
+
+; Negative test: mul(zext x, zext y) may overflow.
+define i32 @mul_may_overflow(i32 %x, i32 %y) {
+; CHECK-LABEL: @mul_may_overflow(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[L:%.*]] = zext i32 [[X:%.*]] to i34
+; CHECK-NEXT:    [[R:%.*]] = zext i32 [[Y:%.*]] to i34
+; CHECK-NEXT:    [[MUL34:%.*]] = mul i34 [[L]], [[R]]
+; CHECK-NEXT:    [[OVERFLOW:%.*]] = icmp ult i34 [[MUL34]], 4294967296
+; CHECK-NEXT:    [[RETVAL:%.*]] = zext i1 [[OVERFLOW]] to i32
+; CHECK-NEXT:    ret i32 [[RETVAL]]
+;
+entry:
+  %l = zext i32 %x to i34
+  %r = zext i32 %y to i34
+  %mul34 = mul i34 %l, %r
+  %overflow = icmp ule i34 %mul34, 4294967295
+  %retval = zext i1 %overflow to i32
+  ret i32 %retval
+}
+
+define i32 @mul_known_nuw(i32 %x, i32 %y) {
+; CHECK-LABEL: @mul_known_nuw(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[UMUL:%.*]] = call { i32, i1 } @llvm.umul.with.overflow.i32(i32 [[X:%.*]], i32 [[Y:%.*]])
+; CHECK-NEXT:    [[TMP0:%.*]] = extractvalue { i32, i1 } [[UMUL]], 1
+; CHECK-NEXT:    [[OVERFLOW:%.*]] = xor i1 [[TMP0]], true
+; CHECK-NEXT:    [[RETVAL:%.*]] = zext i1 [[OVERFLOW]] to i32
+; CHECK-NEXT:    ret i32 [[RETVAL]]
+;
+entry:
+  %l = zext i32 %x to i34
+  %r = zext i32 %y to i34
+  %mul34 = mul nuw i34 %l, %r
+  %overflow = icmp ule i34 %mul34, 4294967295
+  %retval = zext i1 %overflow to i32
+  ret i32 %retval
+}