[GlobalISel] Fix incorrect sign extension when combining G_INTTOPTR and G_PTR_ADD

[pratlucas]

Summary

The GlobalISel combiner currently uses sign extension when manipulating
the LHS constant when combining a sequence of the following sequence of
machine instructions into a single constant:

%0:_(s32) = G_CONSTANT i32 <CONSTANT>
%1:_(p0) = G_INTTOPTR %0:_(s32)
%2:_(s64) = G_CONSTANT i64 <CONSTANT>
%3:_(p0) = G_PTR_ADD %1:_, %2:_(s64)

This causes an issue when the bit width of the first contant and the
target pointer size are different, as G_INTTOPTR has no sign extension
semantics.

This patch fixes this by capture an arbitrary precision in when matching
the constant, allowing the matching function to correctly zero extend
it.

[arsenm]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

2028

Why leave this as int64_t? Why not just leave everything as APInt?

[pratlucas]

Using APInt on match/apply functions' interfaces and merging mir test into existing one.

[arsenm]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

2038

This is a confusing way of stating that inttoptr uses zext

2040

Why sext here?

[pratlucas]

Updating comment.

[pratlucas]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

2040

Using sext here makes sure the pointer arithmetic works properly for negative offsets from G_PTR_ADD. This is also aligned with the behaviour of IRTranslator, which sign-extends the constant offset value (see https://github.com/llvm/llvm-project/blob/main/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp#L1516).
The extension here is not strictly necessary, as the IRTranslator already adjusts the constant size to match the pointer, but it should make things more future-proof without any harm.

llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h

public:
   /// Transform G_ADD (G_PTRTOINT x), y -> G_PTRTOINT (G_PTR_ADD x, y)
   /// Transform G_ADD y, (G_PTRTOINT x) -> G_PTRTOINT (G_PTR_ADD x, y)
   bool matchCombineAddP2IToPtrAdd(MachineInstr &MI,
                                   std::pair<Register, bool> &PtrRegAndCommute);
   void applyCombineAddP2IToPtrAdd(MachineInstr &MI,
                                   std::pair<Register, bool> &PtrRegAndCommute);
 
   // Transform G_PTR_ADD (G_PTRTOINT C1), C2 -> C1 + C2
-  bool matchCombineConstPtrAddToI2P(MachineInstr &MI, int64_t &NewCst);
-  void applyCombineConstPtrAddToI2P(MachineInstr &MI, int64_t &NewCst);
+  bool matchCombineConstPtrAddToI2P(MachineInstr &MI, APInt &NewCst);
+  void applyCombineConstPtrAddToI2P(MachineInstr &MI, APInt &NewCst);
 
   /// Transform anyext(trunc(x)) to x.
   bool matchCombineAnyExtTrunc(MachineInstr &MI, Register &Reg);
   void applyCombineAnyExtTrunc(MachineInstr &MI, Register &Reg);
 
   /// Transform zext(trunc(x)) to x.
   bool matchCombineZextTrunc(MachineInstr &MI, Register &Reg);
 

llvm/include/llvm/CodeGen/GlobalISel/MIPatternMatch.h

 //==------ llvm/CodeGen/GlobalISel/MIPatternMatch.h -------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 /// \file
 /// Contains matchers for matching SSA Machine Instructions.
 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CODEGEN_GLOBALISEL_MIPATTERNMATCH_H
 #define LLVM_CODEGEN_GLOBALISEL_MIPATTERNMATCH_H
 
+#include "llvm/ADT/APInt.h"
 #include "llvm/CodeGen/GlobalISel/Utils.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/IR/InstrTypes.h"
 
 namespace llvm {
 namespace MIPatternMatch {
 
 template <typename Reg, typename Pattern>
template <typename SubPatternT> struct OneNonDBGUse_match {
   }
 };
 
 template <typename SubPat>
 inline OneNonDBGUse_match<SubPat> m_OneNonDBGUse(const SubPat &SP) {
   return SP;
 }
 
-struct ConstantMatch {
-  int64_t &CR;
-  ConstantMatch(int64_t &C) : CR(C) {}
+template <typename ConstT>
+inline Optional<ConstT> matchConstant(Register, const MachineRegisterInfo &);
+
+template <>
+inline Optional<APInt> matchConstant(Register Reg,
+                                     const MachineRegisterInfo &MRI) {
+  return getIConstantVRegVal(Reg, MRI);
+}
+
+template <>
+inline Optional<int64_t> matchConstant(Register Reg,
+                                       const MachineRegisterInfo &MRI) {
+  return getIConstantVRegSExtVal(Reg, MRI);
+}
+
+template <typename ConstT> struct ConstantMatch {
+  ConstT &CR;
+  ConstantMatch(ConstT &C) : CR(C) {}
   bool match(const MachineRegisterInfo &MRI, Register Reg) {
-    if (auto MaybeCst = getIConstantVRegSExtVal(Reg, MRI)) {
+    if (auto MaybeCst = matchConstant<ConstT>(Reg, MRI)) {
       CR = *MaybeCst;
       return true;
     }
     return false;
   }
 };
 
-inline ConstantMatch m_ICst(int64_t &Cst) { return ConstantMatch(Cst); }
+inline ConstantMatch<APInt> m_ICst(APInt &Cst) {
+  return ConstantMatch<APInt>(Cst);
+}
+inline ConstantMatch<int64_t> m_ICst(int64_t &Cst) {
+  return ConstantMatch<int64_t>(Cst);
+}
 
 struct GCstAndRegMatch {
   Optional<ValueAndVReg> &ValReg;
   GCstAndRegMatch(Optional<ValueAndVReg> &ValReg) : ValReg(ValReg) {}
   bool match(const MachineRegisterInfo &MRI, Register Reg) {
     ValReg = getIConstantVRegValWithLookThrough(Reg, MRI);
     return ValReg ? true : false;
   }

llvm/include/llvm/Target/GlobalISel/Combine.td

 def add_p2i_to_ptradd : GICombineRule<
   (defs root:$root, add_p2i_to_ptradd_matchinfo:$info),
   (match (wip_match_opcode G_ADD):$root,
     [{ return Helper.matchCombineAddP2IToPtrAdd(*${root}, ${info}); }]),
   (apply [{ Helper.applyCombineAddP2IToPtrAdd(*${root}, ${info}); }])
 >;
 
 // Fold (ptr_add (int2ptr C1), C2) -> C1 + C2
-def const_ptradd_to_i2p_matchinfo : GIDefMatchData<"int64_t">;
+def const_ptradd_to_i2p_matchinfo : GIDefMatchData<"APInt">;
 def const_ptradd_to_i2p: GICombineRule<
   (defs root:$root, const_ptradd_to_i2p_matchinfo:$info),
   (match (wip_match_opcode G_PTR_ADD):$root,
     [{ return Helper.matchCombineConstPtrAddToI2P(*${root}, ${info}); }]),
   (apply [{ Helper.applyCombineConstPtrAddToI2P(*${root}, ${info}); }])
 >;
 
 // Simplify: (logic_op (op x...), (op y...)) -> (op (logic_op x, y))

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

void CombinerHelper::applyCombineAddP2IToPtrAdd(
 
   Builder.setInstrAndDebugLoc(MI);
   auto PtrAdd = Builder.buildPtrAdd(PtrTy, LHS, RHS);
   Builder.buildPtrToInt(Dst, PtrAdd);
   MI.eraseFromParent();
 }
 
 bool CombinerHelper::matchCombineConstPtrAddToI2P(MachineInstr &MI,
-                                                  int64_t &NewCst) {
+                                                  APInt &NewCst) {
   auto &PtrAdd = cast<GPtrAdd>(MI);
   Register LHS = PtrAdd.getBaseReg();
   Register RHS = PtrAdd.getOffsetReg();
   MachineRegisterInfo &MRI = Builder.getMF().getRegInfo();
 
-  if (auto RHSCst = getIConstantVRegSExtVal(RHS, MRI)) {
-    int64_t Cst;
+  if (auto RHSCst = getIConstantVRegVal(RHS, MRI)) {
+    APInt Cst;
     if (mi_match(LHS, MRI, m_GIntToPtr(m_ICst(Cst)))) {
-      NewCst = Cst + *RHSCst;
+      auto DstTy = MRI.getType(PtrAdd.getReg(0));
+      // G_INTTOPTR uses zero-extension
+      NewCst = Cst.zextOrTrunc(DstTy.getSizeInBits());
+      NewCst += RHSCst->sextOrTrunc(DstTy.getSizeInBits());
       return true;
     }
   }
 
   return false;
 }
 
 void CombinerHelper::applyCombineConstPtrAddToI2P(MachineInstr &MI,
-                                                  int64_t &NewCst) {
+                                                  APInt &NewCst) {
   auto &PtrAdd = cast<GPtrAdd>(MI);
   Register Dst = PtrAdd.getReg(0);
 
   Builder.setInstrAndDebugLoc(MI);
   Builder.buildConstant(Dst, NewCst);
   PtrAdd.eraseFromParent();
 }
 

llvm/test/CodeGen/AArch64/GlobalISel/combine-ptradd-int2ptr.mir

bb.1:
     ; CHECK: [[PTRTOINT:%[0-9]+]]:_(s64) = G_PTRTOINT [[PTR_ADD]](p64)
     ; CHECK: $x0 = COPY [[PTRTOINT]](s64)
     %1:_(s32) = G_CONSTANT i32 42
     %2:_(p64) = COPY $x1
     %3:_(p64) = G_PTR_ADD %2, %1
     %4:_(s64) = G_PTRTOINT %3
     $x0 = COPY %4(s64)
 ...
+---
+name:            test_combine_zero_extend
+body:             |
+  bb.1:
+  liveins: $x0
+
+    ; CHECK-LABEL: name: test_combine_zero_extend
+    ; CHECK: [[C:%[0-9]+]]:_(p0) = G_CONSTANT i64 4291891236
+    ; CHECK-NEXT: $x0 = COPY [[C]](p0)
+    %0:_(s32) = G_CONSTANT i32 -3076096
+    %1:_(p0) = G_INTTOPTR %0:_(s32)
+    %2:_(s64) = G_CONSTANT i64 36
+    %3:_(p0) = G_PTR_ADD %1:_, %2:_(s64)
+    $x0 = COPY %3
+...

llvm/test/CodeGen/AArch64/GlobalISel/inttoptr_add.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64-none-eabi -global-isel -verify-machineinstrs %s -o - | FileCheck %s
+
+define dso_local void @fn() {
+; CHECK-LABEL: fn:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov x8, #4132
+; CHECK-NEXT:    mov w9, #1
+; CHECK-NEXT:    movk x8, #65489, lsl #16
+; CHECK-NEXT:    str w9, [x8]
+; CHECK-NEXT:    ret
+entry:
+  store i32 1, i32* bitcast (i8* getelementptr inbounds (i8, i8* inttoptr (i32 -3076096 to i8*), i64 36) to i32*), align 4
+  ret void
+}