[AMDGPU] Move default initialization of M0 register after the instruction selection

[vpykhtin]

Summary

This allows to avoid problems with glue on extraction of lo16 bit of DS_READ. Note there is a small improvement in ds groupping on pre-gfx9.

There are some cleanup left, its better to made it in separate patch.

[arsenm]

Overall this is probably fine to move. FYI I was hoping to move in a different direction for m0 initialization, where we would stop reserving it and use ordinary copies to initialize it. This would allow us to move/eliminate the custom M0 init optimizations.

Can this remove the SI_INIT_M0 instruction now?

llvm/lib/Target/AMDGPU/DSInstructions.td

55

Could you just set hasPostISelHook to 0 for the special cases?

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

10790

Capitalize

10792

You could just use this in an assert if you disabled the post isel hook for the special cases

10821–10822

Should not rely on the memory operands. isDS() would be more reliable

10825

Make a member function of SITargetLowering and avoid the need to find Subtarget etc.

10828

Check the GDS bit

10851

TII->isDS()?

[vpykhtin]

The SI_INIT_M0 is still used for instructions that I mention in hasNonDefaultM0. Comments say it was introduced to produce S_MOV_B32 m0 so the CSE could join them.

There're a few instructions that set M0 out of intrinsic arguments like V_INTERP, they produce COPY to M0.

llvm/lib/Target/AMDGPU/DSInstructions.td

55

I thought to do it this way, the hook is also used for converting to no-ret atomics, I need to check if there is a case for atomic but not for default M0 init.

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

10792

need to check with intersection on no-ret atomics.

[vpykhtin]

Updated patch. Everything is done except I decided to left readsM0 check for the no-ret atomics case.

[vpykhtin]

llvm/lib/Target/AMDGPU/DSInstructions.td

55

I decided to left readM0 check to distinguish from no-ret atomics usage.

[arsenm]

I just had the realization this may be more appropriate to place in EmitInstrWithCustomInserter, rather than AdjustInstrPostInstrSelection. Since we have to specially treat most of the other special case DS instructions, you'll avoid the need to blacklist more of them. There might be an issue with where the verifier runs between isel and finalize-isel though

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

10851

Shouldn't rely on getNumMemOperands. SelectionDAG has been known to occasionally drop memory operands

[arsenm]

Marking requested changes based on previous comment

llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp

private:
   bool isVGPRImm(const SDNode *N) const;
   bool isUniformLoad(const SDNode *N) const;
   bool isUniformBr(const SDNode *N) const;
 
   MachineSDNode *buildSMovImm64(SDLoc &DL, uint64_t Val, EVT VT) const;
 
   SDNode *glueCopyToOp(SDNode *N, SDValue NewChain, SDValue Glue) const;
   SDNode *glueCopyToM0(SDNode *N, SDValue Val) const;
-  SDNode *glueCopyToM0LDSInit(SDNode *N) const;
 
   const TargetRegisterClass *getOperandRegClass(SDNode *N, unsigned OpNo) const;
   virtual bool SelectADDRVTX_READ(SDValue Addr, SDValue &Base, SDValue &Offset);
   virtual bool SelectADDRIndirect(SDValue Addr, SDValue &Base, SDValue &Offset);
   bool isDSOffsetLegal(SDValue Base, unsigned Offset,
                        unsigned OffsetBits) const;
   bool SelectDS1Addr1Offset(SDValue Ptr, SDValue &Base, SDValue &Offset) const;
   bool SelectDS64Bit4ByteAligned(SDValue Ptr, SDValue &Base, SDValue &Offset0,
const SITargetLowering& Lowering =
     *static_cast<const SITargetLowering*>(getTargetLowering());
 
   assert(N->getOperand(0).getValueType() == MVT::Other && "Expected chain");
 
   SDValue M0 = Lowering.copyToM0(*CurDAG, N->getOperand(0), SDLoc(N), Val);
   return glueCopyToOp(N, M0, M0.getValue(1));
 }
 
-SDNode *AMDGPUDAGToDAGISel::glueCopyToM0LDSInit(SDNode *N) const {
-  unsigned AS = cast<MemSDNode>(N)->getAddressSpace();
-  if (AS == AMDGPUAS::LOCAL_ADDRESS) {
-    if (Subtarget->ldsRequiresM0Init())
-      return glueCopyToM0(N, CurDAG->getTargetConstant(-1, SDLoc(N), MVT::i32));
-  } else if (AS == AMDGPUAS::REGION_ADDRESS) {
-    MachineFunction &MF = CurDAG->getMachineFunction();
-    unsigned Value = MF.getInfo<SIMachineFunctionInfo>()->getGDSSize();
-    return
-        glueCopyToM0(N, CurDAG->getTargetConstant(Value, SDLoc(N), MVT::i32));
-  }
-  return N;
-}
-
 MachineSDNode *AMDGPUDAGToDAGISel::buildSMovImm64(SDLoc &DL, uint64_t Imm,
                                                   EVT VT) const {
   SDNode *Lo = CurDAG->getMachineNode(
       AMDGPU::S_MOV_B32, DL, MVT::i32,
       CurDAG->getTargetConstant(Imm & 0xFFFFFFFF, DL, MVT::i32));
   SDNode *Hi =
       CurDAG->getMachineNode(AMDGPU::S_MOV_B32, DL, MVT::i32,
                              CurDAG->getTargetConstant(Imm >> 32, DL, MVT::i32));

 
 void AMDGPUDAGToDAGISel::Select(SDNode *N) {
   unsigned int Opc = N->getOpcode();
   if (N->isMachineOpcode()) {
     N->setNodeId(-1);
     return;   // Already selected.
   }
 
-  // isa<MemSDNode> almost works but is slightly too permissive for some DS
-  // intrinsics.
-  if (Opc == ISD::LOAD || Opc == ISD::STORE || isa<AtomicSDNode>(N) ||
-      (Opc == AMDGPUISD::ATOMIC_INC || Opc == AMDGPUISD::ATOMIC_DEC ||
-       Opc == ISD::ATOMIC_LOAD_FADD ||
-       Opc == AMDGPUISD::ATOMIC_LOAD_FMIN ||
-       Opc == AMDGPUISD::ATOMIC_LOAD_FMAX)) {
-    N = glueCopyToM0LDSInit(N);
-    SelectCode(N);
-    return;
-  }
-
   switch (Opc) {
   default:
     break;
   // We are selecting i64 ADD here instead of custom lower it during
   // DAG legalization, so we can fold some i64 ADDs used for address
   // calculation into the LOAD and STORE instructions.
   case ISD::ADDC:
   case ISD::ADDE:

llvm/lib/Target/AMDGPU/DSInstructions.td

class DS_Pseudo <string opName, dag outs, dag ins, string asmOps, list<dag> pattern=[]> :
   bits<1> has_offset1 = 1;
 
   bits<1> has_gds = 1;
   bits<1> gdsValue = 0; // if has_gds == 0 set gds to this value
 
   bits<1> has_m0_read = 1;
 
   let Uses = !if(has_m0_read, [M0, EXEC], [EXEC]);
+
+  // initialize M0 in SITargetLowering::AdjustInstrPostInstrSelection
+  let hasPostISelHook = has_m0_read;
 }
 
 class DS_Real <DS_Pseudo ds> :
   InstSI <ds.OutOperandList, ds.InOperandList, ds.Mnemonic # " " # ds.AsmOperands, []>,
   Enc64 {
 
   let isPseudo = 0;
   let isCodeGenOnly = 0;

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

SDNode *SITargetLowering::PostISelFolding(MachineSDNode *Node,
   }
   default:
     break;
   }
 
   return Node;
 }
 
+// these instructions specify M0 value during selection,
+// see AMDGPUDAGToDAGISel::glueCopyToM0, SITargetLowering::copyToM0 usage
+static bool hasNonDefaultM0(unsigned Opcode) {
+  switch(Opcode) {
+  case AMDGPU::DS_APPEND:
+  case AMDGPU::DS_CONSUME:
+  case AMDGPU::DS_GWS_INIT:
+  case AMDGPU::DS_GWS_BARRIER:
+  case AMDGPU::DS_GWS_SEMA_V:
+  case AMDGPU::DS_GWS_SEMA_BR:
+  case AMDGPU::DS_GWS_SEMA_P:
+  case AMDGPU::DS_GWS_SEMA_RELEASE_ALL:
+  case AMDGPU::DS_ORDERED_COUNT:
+    return true;
+  default: break;
+  }
+  return false;
+}
+
+// faster version of readsRegister: M0 is usually the first implicit
+static bool readsM0(MachineInstr &MI) {
+  for (unsigned I = MI.getNumExplicitOperands(), E = MI.getNumOperands();
+       I != E; ++I) {
+    auto &MO = MI.getOperand(I);
+    if (MO.isReg() && MO.isUse() && MO.getReg() == AMDGPU::M0)
+      return true;
+  }
+  return false;
+}
+
+static void initDefaultM0(MachineInstr &MI) {
+  if (!MI.getNumMemOperands())
+    return;
+
+  MachineFunction &MF = *MI.getParent()->getParent();
+  auto &Subtarget = static_cast<const GCNSubtarget&>(MF.getSubtarget());
+  auto AS = (*MI.memoperands_begin())->getAddrSpace();
+  unsigned M0Init;
+  if (AS == AMDGPUAS::REGION_ADDRESS)
+    M0Init = MF.getInfo<SIMachineFunctionInfo>()->getGDSSize();
+  else if (AS == AMDGPUAS::LOCAL_ADDRESS && Subtarget.ldsRequiresM0Init())
+    M0Init = -1;
+  else
+    return;
+
+  if (!readsM0(MI) || hasNonDefaultM0(MI.getOpcode()))
+    return;
+
+  auto &TII = *Subtarget.getInstrInfo();
+  BuildMI(*MI.getParent(), MI, MI.getDebugLoc(),
+          TII.get(AMDGPU::S_MOV_B32), AMDGPU::M0).addImm(M0Init);
+}
+
 /// Assign the register class depending on the number of
 /// bits set in the writemask
 void SITargetLowering::AdjustInstrPostInstrSelection(MachineInstr &MI,
                                                      SDNode *Node) const {
   const SIInstrInfo *TII = getSubtarget()->getInstrInfo();
 
   MachineRegisterInfo &MRI = MI.getParent()->getParent()->getRegInfo();
 
+  initDefaultM0(MI);
+
   if (TII->isVOP3(MI.getOpcode())) {
     // Make sure constant bus requirements are respected.
     TII->legalizeOperandsVOP3(MRI, MI);
 
     // Prefer VGPRs over AGPRs in mAI instructions where possible.
     // This saves a chain-copy of registers and better ballance register
     // use between vgpr and agpr as agpr tuples tend to be big.
     if (const MCOperandInfo *OpInfo = MI.getDesc().OpInfo) {

llvm/test/CodeGen/AMDGPU/ds_read2.ll

define amdgpu_kernel void @ds_read_call_read(i32 addrspace(1)* %out, i32 addrspace(3)* %arg) {
   call void @void_func_void()
   %v1 = load i32, i32 addrspace(3)* %arrayidx1, align 4
   %r = add i32 %v0, %v1
   store i32 %r, i32 addrspace(1)* %out, align 4
   ret void
 }
 
 ; GCN-LABEL: {{^}}ds_read_interp_read:
-; CI: s_mov_b32 m0, -1
-; CI: ds_read_b32
 ; CI: s_mov_b32 m0, s0
 ; CI: v_interp_mov_f32
 ; CI: s_mov_b32 m0, -1
-; CI: ds_read_b32
+; CI: ds_read2_b32 v[0:1], v0 offset1:4
 ; GFX9: ds_read2_b32 v[0:1], v0 offset1:4
 ; GFX9: s_mov_b32 m0, s0
 ; GFX9: v_interp_mov_f32
 define amdgpu_ps <2 x float> @ds_read_interp_read(i32 inreg %prims, float addrspace(3)* %inptr) {
   %v0 = load float, float addrspace(3)* %inptr, align 4
   %intrp = call float @llvm.amdgcn.interp.mov(i32 0, i32 0, i32 0, i32 %prims)
   %ptr1 = getelementptr float, float addrspace(3)* %inptr, i32 4
   %v1 = load float, float addrspace(3)* %ptr1, align 4

llvm/test/CodeGen/AMDGPU/insert-subvector-unused-scratch.ll

 ; RUN: llc -mtriple amdgcn-amd-- -mcpu=bonaire -verify-machineinstrs < %s | FileCheck -check-prefix=GCN %s
 
 ; Before the fix that this test was committed with, this code would leave
 ; an unused stack slot, causing ScratchSize to be non-zero.
 
 ; GCN-LABEL: store_v3i32:
-; GCN:        ds_read_b32
 ; GCN:        ds_read_b64
+; GCN:        ds_read_b32
 ; GCN:        ds_write_b32
 ; GCN:        ds_write_b64
 ; GCN: ScratchSize: 0
 define amdgpu_kernel void @store_v3i32(<3 x i32> addrspace(3)* %out, <3 x i32> %a) nounwind {
   %val = load <3 x i32>, <3 x i32> addrspace(3)* %out
   %val.1 = add <3 x i32> %a, %val
   store <3 x i32> %val.1, <3 x i32> addrspace(3)* %out, align 16
   ret void
 }
 
 ; GCN-LABEL: store_v5i32:
-; GCN:        ds_read_b32
 ; GCN:        ds_read_b128
+; GCN:        ds_read_b32
 ; GCN:        ds_write_b32
 ; GCN:        ds_write_b128
 ; GCN: ScratchSize: 0
 define amdgpu_kernel void @store_v5i32(<5 x i32> addrspace(3)* %out, <5 x i32> %a) nounwind {
   %val = load <5 x i32>, <5 x i32> addrspace(3)* %out
   %val.1 = add <5 x i32> %a, %val
   store <5 x i32> %val.1, <5 x i32> addrspace(3)* %out, align 16
   ret void
 }