llvm-reduce: Fix incorrect cloning of MachineMemOperands

[arsenm]

Summary

There were two problems with directly copying the MMOs from the old
function. The MMOs are owned by the function's Allocator, so need to
be reallocated anyways (surprisingly I didn't notice breakage on
this). Second, the PseudoSourceValues are also allocated per function
and need to be reallocated.

[MatzeB]

LGTM

[arsenm]

717209763e17c10aaddce9db3faac4ecbf1afd29

llvm/test/tools/llvm-reduce/mir/preserve-mem-operands.mir

+# REQUIRES: amdgpu-registered-target
+# RUN: llvm-reduce -simplify-mir -mtriple=amdgcn-amd-amdhsa --test FileCheck --test-arg --check-prefix=CHECK-INTERESTINGNESS --test-arg %s --test-arg --input-file %s -o %t 2> %t.log
+# RUN: FileCheck --match-full-lines --check-prefix=RESULT %s < %t
+
+# CHECK-INTERESTINGNESS: G_LOAD
+# CHECK-INTERESTINGNESS: G_LOAD
+# CHECK-INTERESTINGNESS: G_LOAD
+# CHECK-INTERESTINGNESS: G_LOAD
+# CHECK-INTERESTINGNESS: G_STORE
+# CHECK-INTERESTINGNESS: G_STORE
+# CHECK-INTERESTINGNESS: G_STORE
+# CHECK-INTERESTINGNESS: G_LOAD
+# CHECK-INTERESTINGNESS: G_LOAD
+# CHECK-INTERESTINGNESS: G_LOAD
+# CHECK-INTERESTINGNESS: G_LOAD
+# CHECK-INTERESTINGNESS: G_STORE
+# CHECK-INTERESTINGNESS: G_STORE
+
+
+# RESULT: %{{[0-9]+}}:_(<2 x s16>) = G_LOAD %{{[0-9]+}}(p1) :: (load (<2 x s16>) from %ir.argptr0, align 32, addrspace 1)
+# RESULT: %{{[0-9]+}}:_(<2 x s32>) = G_LOAD %{{[0-9]+}}(p1) :: (load (<2 x s32>) from %ir.argptr1, addrspace 3)
+# RESULT: %{{[0-9]+}}:_(<2 x s32>) = G_LOAD %{{[0-9]+}}(p1) :: (load (<2 x s32>) from %ir.argptr1 + 8, addrspace 3)
+# RESULT: %{{[0-9]+}}:_(<2 x s32>) = G_LOAD %{{[0-9]+}}(p1) :: (load (<2 x s32>) from %ir.argptr1 + 12, align 4, basealign 8, addrspace 3)
+# RESULT: %{{[0-9]+}}:_(p3) = G_IMPLICIT_DEF
+# RESULT: G_STORE %{{[0-9]+}}(<2 x s32>), %{{[0-9]+}}(p3) :: (store (<2 x s32>) into %fixed-stack.0, addrspace 5)
+# RESULT: G_STORE %{{[0-9]+}}(<2 x s32>), %{{[0-9]+}}(p3) :: (store (<2 x s32>) into %stack.0, addrspace 5)
+# RESULT: G_STORE %{{[0-9]+}}(p3), %{{[0-9]+}}(p3) :: (store (p3), addrspace 5)
+# RESULT: %{{[0-9]+}}:_(s32) = G_LOAD %{{[0-9]+}}(p3) :: (load (s32) from call-entry @foo, addrspace 4)
+# RESULT: %{{[0-9]+}}:_(s32) = G_LOAD %{{[0-9]+}}(p1) :: (load (s32) from constant-pool, addrspace 1)
+# RESULT: %{{[0-9]+}}:_(p1) = G_LOAD %{{[0-9]+}}(p3) :: (load (p1) from got, addrspace 4)
+# RESULT: %{{[0-9]+}}:_(p1) = G_LOAD %{{[0-9]+}}(p3) :: (load (p1) from jump-table, addrspace 4)
+# RESULT: G_STORE %{{[0-9]+}}(<3 x s32>), %{{[0-9]+}}(p5) :: (store (<3 x s32>) into stack, align 8, addrspace 5)
+# RESULT: G_STORE %{{[0-9]+}}(<3 x s32>), %{{[0-9]+}}(p5) :: (store (<3 x s32>) into stack + 12, align 4, basealign 8, addrspace 5)
+--- |
+  declare i32 @foo(i32)
+
+  define void @func(<2 x i16> addrspace(1)* %argptr0, <2 x i32> addrspace(3)* %argptr1)  {
+    ret void
+  }
+
+...
+---
+name: func
+tracksRegLiveness: true
+fixedStack:
+  - { id: 0, offset: 16, size: 8, alignment: 4 }
+stack:
+  - { id: 0, size: 4, alignment: 16 }
+body:             |
+  bb.0:
+    S_WAITCNT 0
+    S_NOP 0
+
+    %0:_(p1) = G_IMPLICIT_DEF
+    %1:_(<2 x s16>) = G_LOAD %0 :: (load (<2 x s16>) from %ir.argptr0, align 32, addrspace 1)
+    %2:_(<2 x s32>) = G_ZEXT %1
+    %3:_(<2 x s32>) = G_LOAD %0 :: (load (<2 x s32>) from %ir.argptr1, addrspace 3)
+    %4:_(<2 x s32>) = G_LOAD %0 :: (load (<2 x s32>) from %ir.argptr1 + 8, addrspace 3)
+    %5:_(<2 x s32>) = G_LOAD %0 :: (load (<2 x s32>) from %ir.argptr1 + 12, addrspace 3)
+    %6:_(<2 x s32>) = G_ADD %2, %3
+    %7:_(<2 x s32>) = G_ADD %6, %4
+    %8:_(<2 x s32>) = G_ADD %7, %5
+    %9:_(p3) = G_IMPLICIT_DEF
+    G_STORE %8, %9 :: (store (<2 x s32>) into %fixed-stack.0, addrspace 5)
+    G_STORE %8, %9 :: (store (<2 x s32>) into %stack.0, addrspace 5)
+
+    ; Check address space, no value
+    G_STORE %9, %9 :: (store (p3), addrspace 5)
+
+    %10:_(p0) = G_IMPLICIT_DEF
+    %11:_(s32) = G_LOAD %10 :: (load (s32) from call-entry @foo, addrspace 4)
+    %12:_(s32) = G_LOAD %0 :: (load (s32) from constant-pool, addrspace 1)
+
+    %13:_(p1) = G_LOAD %10 :: (load (p1) from got, addrspace 4)
+
+    %14:_(p1) = G_LOAD %10 :: (load (p1) from jump-table, addrspace 4)
+
+    %15:_(<3 x s32>) = G_IMPLICIT_DEF
+    %16:_(p5) = G_IMPLICIT_DEF
+    G_STORE %15, %16 :: (store (<3 x s32>) into stack, align 8, addrspace 5)
+    G_STORE %15, %16 :: (store (<3 x s32>) into stack + 12, basealign 8, addrspace 5)
+
+    S_ENDPGM 0, implicit %10, implicit %12, implicit %13, implicit %14
+
+...

llvm/tools/llvm-reduce/ReducerWorkItem.cpp

static void cloneFrameInfo(
   }
 
   // FIXME: Needs test, missing MIR serialization.
   if (SrcMFI.hasFunctionContextIndex()) {
     DstMFI.setFunctionContextIndex(SrcMFI.getFunctionContextIndex());
   }
 }
 
+static void cloneMemOperands(MachineInstr &DstMI, MachineInstr &SrcMI,
+                             MachineFunction &SrcMF, MachineFunction &DstMF) {
+  // The new MachineMemOperands should be owned by the new function's
+  // Allocator.
+  PseudoSourceValueManager &PSVMgr = DstMF.getPSVManager();
+
+  // We also need to remap the PseudoSourceValues from the new function's
+  // PseudoSourceValueManager.
+  SmallVector<MachineMemOperand *, 2> NewMMOs;
+  for (MachineMemOperand *OldMMO : SrcMI.memoperands()) {
+    MachinePointerInfo NewPtrInfo(OldMMO->getPointerInfo());
+    if (const PseudoSourceValue *PSV =
+            NewPtrInfo.V.dyn_cast<const PseudoSourceValue *>()) {
+      switch (PSV->kind()) {
+      case PseudoSourceValue::Stack:
+        NewPtrInfo.V = PSVMgr.getStack();
+        break;
+      case PseudoSourceValue::GOT:
+        NewPtrInfo.V = PSVMgr.getGOT();
+        break;
+      case PseudoSourceValue::JumpTable:
+        NewPtrInfo.V = PSVMgr.getJumpTable();
+        break;
+      case PseudoSourceValue::ConstantPool:
+        NewPtrInfo.V = PSVMgr.getConstantPool();
+        break;
+      case PseudoSourceValue::FixedStack:
+        NewPtrInfo.V = PSVMgr.getFixedStack(
+            cast<FixedStackPseudoSourceValue>(PSV)->getFrameIndex());
+        break;
+      case PseudoSourceValue::GlobalValueCallEntry:
+        NewPtrInfo.V = PSVMgr.getGlobalValueCallEntry(
+            cast<GlobalValuePseudoSourceValue>(PSV)->getValue());
+        break;
+      case PseudoSourceValue::ExternalSymbolCallEntry:
+        NewPtrInfo.V = PSVMgr.getExternalSymbolCallEntry(
+            cast<ExternalSymbolPseudoSourceValue>(PSV)->getSymbol());
+        break;
+      case PseudoSourceValue::TargetCustom:
+      default:
+        // FIXME: We have no generic interface for allocating custom PSVs.
+        report_fatal_error("Cloning TargetCustom PSV not handled");
+      }
+    }
+
+    MachineMemOperand *NewMMO = DstMF.getMachineMemOperand(
+        NewPtrInfo, OldMMO->getFlags(), OldMMO->getMemoryType(),
+        OldMMO->getBaseAlign(), OldMMO->getAAInfo(), OldMMO->getRanges(),
+        OldMMO->getSyncScopeID(), OldMMO->getSuccessOrdering(),
+        OldMMO->getFailureOrdering());
+    NewMMOs.push_back(NewMMO);
+  }
+
+  DstMI.setMemRefs(DstMF, NewMMOs);
+}
+
 static std::unique_ptr<MachineFunction> cloneMF(MachineFunction *SrcMF) {
   auto DstMF = std::make_unique<MachineFunction>(
       SrcMF->getFunction(), SrcMF->getTarget(), SrcMF->getSubtarget(),
       SrcMF->getFunctionNumber(), SrcMF->getMMI());
   DenseMap<MachineBasicBlock *, MachineBasicBlock *> Src2DstMBB;
 
   auto *SrcMRI = &SrcMF->getRegInfo();
   auto *DstMRI = &DstMF->getRegInfo();
for (auto &SrcMI : SrcMBB) {
         // Update MBB.
         if (DstMO.isMBB())
           DstMO.setMBB(Src2DstMBB[DstMO.getMBB()]);
         else if (DstMO.isRegMask())
           DstMRI->addPhysRegsUsedFromRegMask(DstMO.getRegMask());
 
         DstMI->addOperand(DstMO);
       }
-      DstMI->setMemRefs(*DstMF, SrcMI.memoperands());
+
+      cloneMemOperands(*DstMI, SrcMI, *SrcMF, *DstMF);
     }
   }
 
   DstMF->setAlignment(SrcMF->getAlignment());
   DstMF->setExposesReturnsTwice(SrcMF->exposesReturnsTwice());
   DstMF->setHasInlineAsm(SrcMF->hasInlineAsm());
   DstMF->setHasWinCFI(SrcMF->hasWinCFI());