[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/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/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) {
   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;
+    APInt Cst;
     if (mi_match(LHS, MRI, m_GIntToPtr(m_ICst(Cst)))) {
-      NewCst = Cst + *RHSCst;
+      NewCst = Cst.getZExtValue() + *RHSCst;
       return true;
     }
   }
 
   return false;
 }
 
 void CombinerHelper::applyCombineConstPtrAddToI2P(MachineInstr &MI,
                                                   int64_t &NewCst) {
   auto &PtrAdd = cast<GPtrAdd>(MI);
   Register Dst = PtrAdd.getReg(0);

llvm/test/CodeGen/AArch64/GlobalISel/combine-inttoptr-add.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -run-pass=aarch64-prelegalizer-combiner -verify-machineinstrs -mtriple aarch64-unknown-unknown %s -o - | FileCheck %s
+---
+name:            test_combine_inttoptr_add
+body:             |
+  bb.1:
+  liveins: $x0
+
+    ; CHECK-LABEL: name: test_combine_inttoptr_add
+    ; 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
+}