[MIR] Add support for addrspace in MIR

[thegameg]

Summary

Add support for printing / parsing the addrspace of a MachineMemOperand.

PR35970

[MatzeB]

LGTM

(+CC AMDGPU so they can veto if they don't like the syntax)

test/CodeGen/MIR/AArch64/addrspace-memoperands.mir

27

Can you add a line that addrspace 0 is correctly parsed (and not printed)

[thegameg]

Add test case with addrspace 0.

[thegameg]

Update some more failing tests.

lib/CodeGen/MIRParser/MILexer.h

enum TokenKind {
     kw_x86_fp80,
     kw_fp128,
     kw_ppc_fp128,
     kw_target_flags,
     kw_volatile,
     kw_non_temporal,
     kw_invariant,
     kw_align,
+    kw_addrspace,
     kw_stack,
     kw_got,
     kw_jump_table,
     kw_constant_pool,
     kw_call_entry,
     kw_liveout,
     kw_address_taken,
     kw_landing_pad,

lib/CodeGen/MIRParser/MILexer.cpp

return StringSwitch<MIToken::TokenKind>(Identifier)
       .Case("fp128", MIToken::kw_fp128)
       .Case("ppc_fp128", MIToken::kw_ppc_fp128)
       .Case("target-flags", MIToken::kw_target_flags)
       .Case("volatile", MIToken::kw_volatile)
       .Case("non-temporal", MIToken::kw_non_temporal)
       .Case("dereferenceable", MIToken::kw_dereferenceable)
       .Case("invariant", MIToken::kw_invariant)
       .Case("align", MIToken::kw_align)
+      .Case("addrspace", MIToken::kw_addrspace)
       .Case("stack", MIToken::kw_stack)
       .Case("got", MIToken::kw_got)
       .Case("jump-table", MIToken::kw_jump_table)
       .Case("constant-pool", MIToken::kw_constant_pool)
       .Case("call-entry", MIToken::kw_call_entry)
       .Case("liveout", MIToken::kw_liveout)
       .Case("address-taken", MIToken::kw_address_taken)
       .Case("landing-pad", MIToken::kw_landing_pad)

lib/CodeGen/MIRParser/MIParser.cpp

public:
   bool parseCustomRegisterMaskOperand(MachineOperand &Dest);
   bool parseLiveoutRegisterMaskOperand(MachineOperand &Dest);
   bool parseMachineOperand(MachineOperand &Dest,
                            Optional<unsigned> &TiedDefIdx);
   bool parseMachineOperandAndTargetFlags(MachineOperand &Dest,
                                          Optional<unsigned> &TiedDefIdx);
   bool parseOffset(int64_t &Offset);
   bool parseAlignment(unsigned &Alignment);
+  bool parseAddrspace(unsigned &Addrspace);
   bool parseOperandsOffset(MachineOperand &Op);
   bool parseIRValue(const Value *&V);
   bool parseMemoryOperandFlag(MachineMemOperand::Flags &Flags);
   bool parseMemoryPseudoSourceValue(const PseudoSourceValue *&PSV);
   bool parseMachinePointerInfo(MachinePointerInfo &Dest);
   bool parseOptionalScope(LLVMContext &Context, SyncScope::ID &SSID);
   bool parseOptionalAtomicOrdering(AtomicOrdering &Order);
   bool parseMachineMemoryOperand(MachineMemOperand *&Dest);
bool MIParser::parseAlignment(unsigned &Alignment) {
   if (Token.isNot(MIToken::IntegerLiteral) || Token.integerValue().isSigned())
     return error("expected an integer literal after 'align'");
   if (getUnsigned(Alignment))
     return true;
   lex();
   return false;
 }
 
+bool MIParser::parseAddrspace(unsigned &Addrspace) {
+  assert(Token.is(MIToken::kw_addrspace));
+  lex();
+  if (Token.isNot(MIToken::IntegerLiteral) || Token.integerValue().isSigned())
+    return error("expected an integer literal after 'addrspace'");
+  if (getUnsigned(Addrspace))
+    return true;
+  lex();
+  return false;
+}
+
 bool MIParser::parseOperandsOffset(MachineOperand &Op) {
   int64_t Offset = 0;
   if (parseOffset(Offset))
     return true;
   Op.setOffset(Offset);
   return false;
 }
 
bool MIParser::parseMachineMemoryOperand(MachineMemOperand *&Dest) {
   AAMDNodes AAInfo;
   MDNode *Range = nullptr;
   while (consumeIfPresent(MIToken::comma)) {
     switch (Token.kind()) {
     case MIToken::kw_align:
       if (parseAlignment(BaseAlignment))
         return true;
       break;
+    case MIToken::kw_addrspace:
+      if (parseAddrspace(Ptr.AddrSpace))
+        return true;
+      break;
     case MIToken::md_tbaa:
       lex();
       if (parseMDNode(AAInfo.TBAA))
         return true;
       break;
     case MIToken::md_alias_scope:
       lex();
       if (parseMDNode(AAInfo.Scope))

lib/CodeGen/MIRPrinter.cpp

void MIPrinter::print(const LLVMContext &Context, const TargetInstrInfo &TII,
   if (AAInfo.NoAlias) {
     OS << ", !noalias ";
     AAInfo.NoAlias->printAsOperand(OS, MST);
   }
   if (Op.getRanges()) {
     OS << ", !range ";
     Op.getRanges()->printAsOperand(OS, MST);
   }
+  if (unsigned AS = Op.getAddrSpace())
+    OS << ", addrspace " << AS;
   OS << ')';
 }
 
 void MIPrinter::printSyncScope(const LLVMContext &Context, SyncScope::ID SSID) {
   switch (SSID) {
   case SyncScope::System: {
     break;
   }

test/CodeGen/AArch64/GlobalISel/arm64-irtranslator.ll

define void @trunc(i64 %a) {
   %res2 = trunc <4 x i32> %vec to <4 x i16>
   ret void
 }
 
 ; CHECK-LABEL: name: load
 ; CHECK: [[ADDR:%[0-9]+]]:_(p0) = COPY %x0
 ; CHECK: [[ADDR42:%[0-9]+]]:_(p42) = COPY %x1
 ; CHECK: [[VAL1:%[0-9]+]]:_(s64) = G_LOAD [[ADDR]](p0) :: (load 8 from %ir.addr, align 16)
-; CHECK: [[VAL2:%[0-9]+]]:_(s64) = G_LOAD [[ADDR42]](p42) :: (load 8 from %ir.addr42)
+; CHECK: [[VAL2:%[0-9]+]]:_(s64) = G_LOAD [[ADDR42]](p42) :: (load 8 from %ir.addr42, addrspace 42)
 ; CHECK: [[SUM2:%.*]]:_(s64) = G_ADD [[VAL1]], [[VAL2]]
 ; CHECK: [[VAL3:%[0-9]+]]:_(s64) = G_LOAD [[ADDR]](p0) :: (volatile load 8 from %ir.addr)
 ; CHECK: [[SUM3:%[0-9]+]]:_(s64) = G_ADD [[SUM2]], [[VAL3]]
 ; CHECK: %x0 = COPY [[SUM3]]
 ; CHECK: RET_ReallyLR implicit %x0
 define i64 @load(i64* %addr, i64 addrspace(42)* %addr42) {
   %val1 = load i64, i64* %addr, align 16
 
   %val2 = load i64, i64 addrspace(42)* %addr42
   %sum2 = add i64 %val1, %val2
 
   %val3 = load volatile i64, i64* %addr
   %sum3 = add i64 %sum2, %val3
   ret i64 %sum3
 }
 
 ; CHECK-LABEL: name: store
 ; CHECK: [[ADDR:%[0-9]+]]:_(p0) = COPY %x0
 ; CHECK: [[ADDR42:%[0-9]+]]:_(p42) = COPY %x1
 ; CHECK: [[VAL1:%[0-9]+]]:_(s64) = COPY %x2
 ; CHECK: [[VAL2:%[0-9]+]]:_(s64) = COPY %x3
 ; CHECK: G_STORE [[VAL1]](s64), [[ADDR]](p0) :: (store 8 into %ir.addr, align 16)
-; CHECK: G_STORE [[VAL2]](s64), [[ADDR42]](p42) :: (store 8 into %ir.addr42)
+; CHECK: G_STORE [[VAL2]](s64), [[ADDR42]](p42) :: (store 8 into %ir.addr42, addrspace 42)
 ; CHECK: G_STORE [[VAL1]](s64), [[ADDR]](p0) :: (volatile store 8 into %ir.addr)
 ; CHECK: RET_ReallyLR
 define void @store(i64* %addr, i64 addrspace(42)* %addr42, i64 %val1, i64 %val2) {
   store i64 %val1, i64* %addr, align 16
   store i64 %val2, i64 addrspace(42)* %addr42
   store volatile i64 %val1, i64* %addr
   %sum = add i64 %val1, %val2
   ret void

test/CodeGen/AMDGPU/optimize-if-exec-masking.mir

bb.2.end:
     %sgpr2 = S_MOV_B32 -1
     BUFFER_STORE_DWORD_OFFSET killed %vgpr0, %sgpr0_sgpr1_sgpr2_sgpr3, 0, 0, 0, 0, 0, implicit %exec :: (store 4 into `i32 addrspace(1)* undef`)
     S_ENDPGM
 
 ...
 ---
 # CHECK-LABEL: name: optimize_if_and_saveexec_xor_valu_middle
 # CHECK: %sgpr2_sgpr3 = S_AND_B64 %sgpr0_sgpr1, killed %vcc, implicit-def %scc
-# CHECK-NEXT: BUFFER_STORE_DWORD_OFFSET %vgpr0, undef %sgpr0_sgpr1_sgpr2_sgpr3, 0, 0, 0, 0, 0, implicit %exec :: (store 4 into `i32 addrspace(1)* undef`)
+# CHECK-NEXT: BUFFER_STORE_DWORD_OFFSET %vgpr0, undef %sgpr0_sgpr1_sgpr2_sgpr3, 0, 0, 0, 0, 0, implicit %exec :: (store 4 into `i32 addrspace(1)* undef`, addrspace 1)
 # CHECK-NEXT: %sgpr0_sgpr1 = S_XOR_B64 %sgpr2_sgpr3, killed %sgpr0_sgpr1, implicit-def %scc
 # CHECK-NEXT: %exec = COPY killed %sgpr2_sgpr3
 # CHECK-NEXT: SI_MASK_BRANCH
 name:            optimize_if_and_saveexec_xor_valu_middle
 alignment:       0
 exposesReturnsTwice: false
 legalized:       false
 regBankSelected: false

test/CodeGen/AMDGPU/syncscopes.ll

 ; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx803 -stop-before=si-debugger-insert-nops < %s | FileCheck --check-prefix=GCN %s
 
 ; GCN-LABEL: name: syncscopes
-; GCN: FLAT_STORE_DWORD killed renamable %vgpr1_vgpr2, killed renamable %vgpr0, 0, 0, 0, implicit %exec, implicit %flat_scr :: (volatile store syncscope("agent") seq_cst 4 into %ir.agent_out)
-; GCN: FLAT_STORE_DWORD killed renamable %vgpr4_vgpr5, killed renamable %vgpr3, 0, 0, 0, implicit %exec, implicit %flat_scr :: (volatile store syncscope("workgroup") seq_cst 4 into %ir.workgroup_out)
-; GCN: FLAT_STORE_DWORD killed renamable %vgpr7_vgpr8, killed renamable %vgpr6, 0, 0, 0, implicit %exec, implicit %flat_scr :: (volatile store syncscope("wavefront") seq_cst 4 into %ir.wavefront_out)
+; GCN: FLAT_STORE_DWORD killed renamable %vgpr1_vgpr2, killed renamable %vgpr0, 0, 0, 0, implicit %exec, implicit %flat_scr :: (volatile store syncscope("agent") seq_cst 4 into %ir.agent_out, addrspace 4)
+; GCN: FLAT_STORE_DWORD killed renamable %vgpr4_vgpr5, killed renamable %vgpr3, 0, 0, 0, implicit %exec, implicit %flat_scr :: (volatile store syncscope("workgroup") seq_cst 4 into %ir.workgroup_out, addrspace 4)
+; GCN: FLAT_STORE_DWORD killed renamable %vgpr7_vgpr8, killed renamable %vgpr6, 0, 0, 0, implicit %exec, implicit %flat_scr :: (volatile store syncscope("wavefront") seq_cst 4 into %ir.wavefront_out, addrspace 4)
 define void @syncscopes(
     i32 %agent,
     i32 addrspace(4)* %agent_out,
     i32 %workgroup,
     i32 addrspace(4)* %workgroup_out,
     i32 %wavefront,
     i32 addrspace(4)* %wavefront_out) {
 entry:
   store atomic i32 %agent, i32 addrspace(4)* %agent_out syncscope("agent") seq_cst, align 4
   store atomic i32 %workgroup, i32 addrspace(4)* %workgroup_out syncscope("workgroup") seq_cst, align 4
   store atomic i32 %wavefront, i32 addrspace(4)* %wavefront_out syncscope("wavefront") seq_cst, align 4
   ret void
 }

test/CodeGen/MIR/AArch64/addrspace-memoperands.mir

+# RUN: llc -mtriple=aarch64-none-linux-gnu -run-pass none -o - %s | FileCheck %s
+
+--- |
+
+  define void @addrspace_memoperands() {
+    ret void
+  }
+
+...
+---
+name:            addrspace_memoperands
+body: |
+  bb.0:
+
+    ; CHECK-LABEL: name: addrspace_memoperands
+    ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY %x0
+    ; CHECK: [[LOAD:%[0-9]+]]:_(s64) = G_LOAD [[COPY]](p0) :: (load 8, addrspace 1)
+    ; CHECK: [[LOAD1:%[0-9]+]]:_(s32) = G_LOAD [[COPY]](p0) :: (load 4, align 2, addrspace 3)
+    ; CHECK: G_STORE [[LOAD]](s64), [[COPY]](p0) :: (store 8, addrspace 1)
+    ; CHECK: G_STORE [[LOAD1]](s32), [[COPY]](p0) :: (store 4, align 2, addrspace 3)
+    ; CHECK: G_STORE [[LOAD1]](s32), [[COPY]](p0) :: (store 4)
+    ; CHECK: RET_ReallyLR
+    %0:_(p0) = COPY %x0
+    %1:_(s64) = G_LOAD %0(p0) :: (load 8, addrspace 1)
+    %2:_(s32) = G_LOAD %0(p0) :: (load 4, align 2, addrspace 3)
+    G_STORE %1(s64), %0(p0) :: (store 8, addrspace 1)
+    G_STORE %2(s32), %0(p0) :: (store 4, align 2, addrspace 3)
+    ; addrspace 0 is accepted by the parser but not printed
+    G_STORE %2(s32), %0(p0) :: (store 4, addrspace 0)
+    RET_ReallyLR
+...

test/CodeGen/MIR/AMDGPU/syncscopes.mir

--- |
   
   attributes #0 = { "target-cpu"="gfx803" }
   attributes #1 = { convergent nounwind }
   attributes #2 = { convergent nounwind readnone }
 
   !0 = !{i32 1}
 
 # GCN-LABEL: name: syncscopes
-# GCN: FLAT_STORE_DWORD killed %vgpr0_vgpr1, killed %vgpr2, 0, -1, 0, implicit %exec, implicit %flat_scr :: (volatile non-temporal store syncscope("agent") seq_cst 4 into %ir.agent_out)
-# GCN: FLAT_STORE_DWORD killed %vgpr0_vgpr1, killed %vgpr2, 0, -1, 0, implicit %exec, implicit %flat_scr :: (volatile non-temporal store syncscope("workgroup") seq_cst 4 into %ir.workgroup_out)
-# GCN: FLAT_STORE_DWORD killed %vgpr0_vgpr1, killed %vgpr2, 0, -1, 0, implicit %exec, implicit %flat_scr :: (volatile non-temporal store syncscope("wavefront") seq_cst 4 into %ir.wavefront_out)
+# GCN: FLAT_STORE_DWORD killed %vgpr0_vgpr1, killed %vgpr2, 0, -1, 0, implicit %exec, implicit %flat_scr :: (volatile non-temporal store syncscope("agent") seq_cst 4 into %ir.agent_out, addrspace 4)
+# GCN: FLAT_STORE_DWORD killed %vgpr0_vgpr1, killed %vgpr2, 0, -1, 0, implicit %exec, implicit %flat_scr :: (volatile non-temporal store syncscope("workgroup") seq_cst 4 into %ir.workgroup_out, addrspace 4)
+# GCN: FLAT_STORE_DWORD killed %vgpr0_vgpr1, killed %vgpr2, 0, -1, 0, implicit %exec, implicit %flat_scr :: (volatile non-temporal store syncscope("wavefront") seq_cst 4 into %ir.wavefront_out, addrspace 4)
 ...
 ---
 name:            syncscopes
 alignment:       0
 exposesReturnsTwice: false
 legalized:       false
 regBankSelected: false
 selected:        false