TableGen: Add address space to matchers

[arsenm]

Summary

Currently AMDGPU uses a CodePatPred to check address spaces from the
MachineMemOperand. Introduce a new first class property so that the
existing patterns can be easily modified to uses the new generated
predicate, which will also be handled for GlobalISel.

I would prefer these to match against the pointer type of the
instruction, but that would be difficult to get working with
SelectionDAG compatbility. This is much easier for now and will avoid
a painful tablegen rewrite for all the loads and stores.

I'm also not sure if there's a better way to encode multiple address
spaces in the table, rather than putting the number to expect.

[dsanders]

LGTM

I'm also not sure if there's a better way to encode multiple address spaces in the table, rather than putting the number to expect.

It amounts to the same thing but they could be LLT's like the ones defined in llvm/Target/GlobalISel/Target.td. If each target defines it's own pX's then it does at least constrain it to the expected address spaces.

test/TableGen/address-space-patfrags.td

20

Regarding the FIXME. The importer currently operates after PatFrag expansion which is too late to infer it. I expect the code that does the expansion could infer it but I haven't dug into that.

utils/TableGen/GlobalISelEmitter.cpp

245–248

Other places in this file achieve this with:

StringRef Separator = "";
for (...) {
  OS << Separator << ...;
  Separator = ", ";
}

I think it ends up a bit neater but I don't mind either way

[arsenm]

r366128

include/llvm/CodeGen/GlobalISel/InstructionSelector.h

enum {
   GIM_CheckAtomicOrdering,
   GIM_CheckAtomicOrderingOrStrongerThan,
   GIM_CheckAtomicOrderingWeakerThan,
   /// Check the size of the memory access for the given machine memory operand.
   /// - InsnID - Instruction ID
   /// - MMOIdx - MMO index
   /// - Size - The size in bytes of the memory access
   GIM_CheckMemorySizeEqualTo,
+
+  /// Check the address space of the memory access for the given machine memory
+  /// operand.
+  /// - InsnID - Instruction ID
+  /// - MMOIdx - MMO index
+  /// - NumAddrSpace - Number of valid address spaces
+  /// - AddrSpaceN - An allowed space of the memory access
+  /// - AddrSpaceN+1 ...
+  GIM_CheckMemoryAddressSpace,
+
   /// Check the size of the memory access for the given machine memory operand
   /// against the size of an operand.
   /// - InsnID - Instruction ID
   /// - MMOIdx - MMO index
   /// - OpIdx - The operand index to compare the MMO against
   GIM_CheckMemorySizeEqualToLLT,
   GIM_CheckMemorySizeLessThanLLT,
   GIM_CheckMemorySizeGreaterThanLLT,

include/llvm/CodeGen/GlobalISel/InstructionSelectorImpl.h

case GIM_CheckAtomicOrderingWeakerThan: {
           return false;
 
       for (const auto &MMO : State.MIs[InsnID]->memoperands())
         if (!isStrongerThan(Ordering, MMO->getOrdering()))
           if (handleReject() == RejectAndGiveUp)
             return false;
       break;
     }
+    case GIM_CheckMemoryAddressSpace: {
+      int64_t InsnID = MatchTable[CurrentIdx++];
+      int64_t MMOIdx = MatchTable[CurrentIdx++];
+      // This accepts a list of possible address spaces.
+      const int NumAddrSpace = MatchTable[CurrentIdx++];
+
+      if (State.MIs[InsnID]->getNumMemOperands() <= MMOIdx) {
+        if (handleReject() == RejectAndGiveUp)
+          return false;
+        break;
+      }
+
+      // Need to still jump to the end of the list of address spaces if we find
+      // a match earlier.
+      const uint64_t LastIdx = CurrentIdx + NumAddrSpace;
+
+      const MachineMemOperand *MMO
+        = *(State.MIs[InsnID]->memoperands_begin() + MMOIdx);
+      const unsigned MMOAddrSpace = MMO->getAddrSpace();
+
+      bool Success = false;
+      for (int I = 0; I != NumAddrSpace; ++I) {
+        unsigned AddrSpace = MatchTable[CurrentIdx++];
+        DEBUG_WITH_TYPE(
+          TgtInstructionSelector::getName(),
+          dbgs() << "addrspace(" << MMOAddrSpace << ") vs "
+                 << AddrSpace << '\n');
+
+        if (AddrSpace == MMOAddrSpace) {
+          Success = true;
+          break;
+        }
+      }
+
+      CurrentIdx = LastIdx;
+      if (!Success && handleReject() == RejectAndGiveUp)
+        return false;
+      break;
+    }
     case GIM_CheckMemorySizeEqualTo: {
       int64_t InsnID = MatchTable[CurrentIdx++];
       int64_t MMOIdx = MatchTable[CurrentIdx++];
       uint64_t Size = MatchTable[CurrentIdx++];
 
       DEBUG_WITH_TYPE(TgtInstructionSelector::getName(),
                       dbgs() << CurrentIdx
                              << ": GIM_CheckMemorySizeEqual(MIs[" << InsnID

include/llvm/Target/TargetSelectionDAG.td

class PatFrags<dag ops, list<dag> frags, code pred = [{}],
   // cast<LoadSDNode>(N)->getExtensionType() == ISD::SEXTLOAD;
   bit IsSignExtLoad = ?;
   // cast<LoadSDNode>(N)->getExtensionType() == ISD::ZEXTLOAD;
   bit IsZeroExtLoad = ?;
   // !cast<StoreSDNode>(N)->isTruncatingStore();
   // cast<StoreSDNode>(N)->isTruncatingStore();
   bit IsTruncStore = ?;
 
+  // cast<MemSDNode>(N)->getAddressSpace() ==
+  // If this empty, accept any address space.
+  list<int> AddressSpaces = ?;
+
   // cast<AtomicSDNode>(N)->getOrdering() == AtomicOrdering::Monotonic
   bit IsAtomicOrderingMonotonic = ?;
   // cast<AtomicSDNode>(N)->getOrdering() == AtomicOrdering::Acquire
   bit IsAtomicOrderingAcquire = ?;
   // cast<AtomicSDNode>(N)->getOrdering() == AtomicOrdering::Release
   bit IsAtomicOrderingRelease = ?;
   // cast<AtomicSDNode>(N)->getOrdering() == AtomicOrdering::AcquireRelease
   bit IsAtomicOrderingAcquireRelease = ?;
class PatFrags<dag ops, list<dag> frags, code pred = [{}],
   bit IsAtomicOrderingReleaseOrStronger = ?;
 
   // cast<LoadSDNode>(N)->getMemoryVT() == MVT::<VT>;
   // cast<StoreSDNode>(N)->getMemoryVT() == MVT::<VT>;
   ValueType MemoryVT = ?;
   // cast<LoadSDNode>(N)->getMemoryVT().getScalarType() == MVT::<VT>;
   // cast<StoreSDNode>(N)->getMemoryVT().getScalarType() == MVT::<VT>;
   ValueType ScalarMemoryVT = ?;
+
+  // TODO: Add alignment
 }
 
 // PatFrag - A version of PatFrags matching only a single fragment.
 class PatFrag<dag ops, dag frag, code pred = [{}],
               SDNodeXForm xform = NOOP_SDNodeXForm>
   : PatFrags<ops, [frag], pred, xform>;
 
 // OutPatFrag is a pattern fragment that is used as part of an output pattern

test/TableGen/address-space-patfrags.td

+// RUN: llvm-tblgen -gen-dag-isel -I %p/../../include %s 2>&1 | FileCheck -check-prefix=SDAG %s
+// RUN: llvm-tblgen -gen-global-isel -optimize-match-table=false -I %p/../../include %s -o - < %s | FileCheck -check-prefix=GISEL %s
+
+include "llvm/Target/Target.td"
+
+def TestTargetInstrInfo : InstrInfo;
+
+
+def TestTarget : Target {
+  let InstructionSet = TestTargetInstrInfo;
+}
+
+def R0 : Register<"r0"> { let Namespace = "MyTarget"; }
+def GPR32 : RegisterClass<"MyTarget", [i32], 32, (add R0)>;
+
+
+// With one address space
+def pat_frag_a : PatFrag <(ops node:$ptr), (load node:$ptr), [{}]> {
+  let AddressSpaces = [ 999 ];
+  let IsLoad = 1; // FIXME: Can this be inferred?
+  let MemoryVT = i32;
+}
+
+// With multiple address spaces
+def pat_frag_b : PatFrag <(ops node:$ptr), (load node:$ptr), [{}]> {
+  let AddressSpaces = [ 123, 455 ];
+  let IsLoad = 1; // FIXME: Can this be inferred?
+  let MemoryVT = i32;
+}
+
+def inst_a : Instruction {
+  let OutOperandList = (outs GPR32:$dst);
+  let InOperandList = (ins GPR32:$src);
+}
+
+def inst_b : Instruction {
+  let OutOperandList = (outs GPR32:$dst);
+  let InOperandList = (ins GPR32:$src);
+}
+
+// SDAG: case 2: {
+// SDAG: // Predicate_pat_frag_a
+// SDAG-NEXT: SDNode *N = Node;
+// SDAG-NEXT: (void)N;
+// SDAG-NEXT: unsigned AddrSpace = cast<MemSDNode>(N)->getAddressSpace();
+
+// SDAG-NEXT: if (AddrSpace != 999)
+// SDAG-NEXT: return false;
+// SDAG-NEXT: if (cast<MemSDNode>(N)->getMemoryVT() != MVT::i32) return false;
+// SDAG-NEXT: return true;
+
+// GISEL: GIM_Try, /*On fail goto*//*Label 0*/ 47, // Rule ID 0 //
+// GISEL-NEXT: GIM_CheckNumOperands, /*MI*/0, /*Expected*/2,
+// GISEL-NEXT: GIM_CheckOpcode, /*MI*/0, TargetOpcode::G_LOAD,
+// GISEL-NEXT: GIM_CheckMemorySizeEqualToLLT, /*MI*/0, /*MMO*/0, /*OpIdx*/0,
+// GISEL-NEXT: GIM_CheckMemoryAddressSpace, /*MI*/0, /*MMO*/0, /*NumAddrSpace*/1, /*AddrSpace*/999,
+// GISEL-NEXT: GIM_CheckMemorySizeEqualTo, /*MI*/0, /*MMO*/0, /*Size*/4,
+// GISEL-NEXT: GIM_CheckAtomicOrdering, /*MI*/0, /*Order*/(int64_t)AtomicOrdering::NotAtomic,
+def : Pat <
+  (pat_frag_a GPR32:$src),
+  (inst_a GPR32:$src)
+>;
+
+// SDAG: case 3: {
+// SDAG-NEXT: // Predicate_pat_frag_b
+// SDAG-NEXT: SDNode *N = Node;
+// SDAG-NEXT: (void)N;
+// SDAG-NEXT: unsigned AddrSpace = cast<MemSDNode>(N)->getAddressSpace();
+// SDAG-NEXT: if (AddrSpace != 123 && AddrSpace != 455)
+// SDAG-NEXT: return false;
+// SDAG-NEXT: if (cast<MemSDNode>(N)->getMemoryVT() != MVT::i32) return false;
+// SDAG-NEXT: return true;
+
+
+// GISEL: GIM_Try, /*On fail goto*//*Label 1*/ 95, // Rule ID 1 //
+// GISEL-NEXT: GIM_CheckNumOperands, /*MI*/0, /*Expected*/2,
+// GISEL-NEXT: GIM_CheckOpcode, /*MI*/0, TargetOpcode::G_LOAD,
+// GISEL-NEXT: GIM_CheckMemorySizeEqualToLLT, /*MI*/0, /*MMO*/0, /*OpIdx*/0,
+// GISEL-NEXT: GIM_CheckMemoryAddressSpace, /*MI*/0, /*MMO*/0, /*NumAddrSpace*/2, /*AddrSpace*/123, /*AddrSpace*/455,
+// GISEL-NEXT: GIM_CheckMemorySizeEqualTo, /*MI*/0, /*MMO*/0, /*Size*/4,
+// GISEL-NEXT: GIM_CheckAtomicOrdering, /*MI*/0, /*Order*/(int64_t)AtomicOrdering::NotAtomic,
+def : Pat <
+  (pat_frag_b GPR32:$src),
+  (inst_b GPR32:$src)
+>;

utils/TableGen/CodeGenDAGPatterns.h

public:
   /// If non-null, indicates that this predicate is a predefined memory VT
   /// predicate for a load/store and returns the ValueType record for the memory VT.
   Record *getMemoryVT() const;
   /// If non-null, indicates that this predicate is a predefined memory VT
   /// predicate (checking only the scalar type) for load/store and returns the
   /// ValueType record for the memory VT.
   Record *getScalarMemoryVT() const;
 
+  ListInit *getAddressSpaces() const;
+
   // If true, indicates that GlobalISel-based C++ code was supplied.
   bool hasGISelPredicateCode() const;
   std::string getGISelPredicateCode() const;
 
 private:
   bool hasPredCode() const;
   bool hasImmCode() const;
   std::string getPredCode() const;

utils/TableGen/CodeGenDAGPatterns.cpp

if (isAtomicOrderingReleaseOrStronger())
       PrintFatalError(getOrigPatFragRecord()->getRecord()->getLoc(),
                       "IsAtomicOrderingReleaseOrStronger requires IsAtomic");
     if (isAtomicOrderingWeakerThanAcquire())
       PrintFatalError(getOrigPatFragRecord()->getRecord()->getLoc(),
                       "IsAtomicOrderingWeakerThanAcquire requires IsAtomic");
   }
 
   if (isLoad() || isStore() || isAtomic()) {
-    StringRef SDNodeName =
-        isLoad() ? "LoadSDNode" : isStore() ? "StoreSDNode" : "AtomicSDNode";
+    if (ListInit *AddressSpaces = getAddressSpaces()) {
+      Code += "unsigned AddrSpace = cast<MemSDNode>(N)->getAddressSpace();\n"
+        " if (";
+
+      bool First = true;
+      for (Init *Val : AddressSpaces->getValues()) {
+        if (First)
+          First = false;
+        else
+          Code += " && ";
+
+        IntInit *IntVal = dyn_cast<IntInit>(Val);
+        if (!IntVal) {
+          PrintFatalError(getOrigPatFragRecord()->getRecord()->getLoc(),
+                          "AddressSpaces element must be integer");
+        }
+
+        Code += "AddrSpace != " + utostr(IntVal->getValue());
+      }
+
+      Code += ")\nreturn false;\n";
+    }
 
     Record *MemoryVT = getMemoryVT();
 
     if (MemoryVT)
-      Code += ("if (cast<" + SDNodeName + ">(N)->getMemoryVT() != MVT::" +
+      Code += ("if (cast<MemSDNode>(N)->getMemoryVT() != MVT::" +
                MemoryVT->getName() + ") return false;\n")
                   .str();
   }
 
   if (isAtomic() && isAtomicOrderingMonotonic())
     Code += "if (cast<AtomicSDNode>(N)->getOrdering() != "
             "AtomicOrdering::Monotonic) return false;\n";
   if (isAtomic() && isAtomicOrderingAcquire())
bool TreePredicateFn::isAtomicOrderingWeakerThanRelease() const {
   return isPredefinedPredicateEqualTo("IsAtomicOrderingReleaseOrStronger", false);
 }
 Record *TreePredicateFn::getMemoryVT() const {
   Record *R = getOrigPatFragRecord()->getRecord();
   if (R->isValueUnset("MemoryVT"))
     return nullptr;
   return R->getValueAsDef("MemoryVT");
 }
+
+ListInit *TreePredicateFn::getAddressSpaces() const {
+  Record *R = getOrigPatFragRecord()->getRecord();
+  if (R->isValueUnset("AddressSpaces"))
+    return nullptr;
+  return R->getValueAsListInit("AddressSpaces");
+}
+
 Record *TreePredicateFn::getScalarMemoryVT() const {
   Record *R = getOrigPatFragRecord()->getRecord();
   if (R->isValueUnset("ScalarMemoryVT"))
     return nullptr;
   return R->getValueAsDef("ScalarMemoryVT");
 }
 bool TreePredicateFn::hasGISelPredicateCode() const {
   return !PatFragRec->getRecord()

utils/TableGen/GlobalISelEmitter.cpp

for (const TreePredicateCall &Call : N->getPredicateCalls()) {
     if (P.isTruncStore())
       Explanation += " truncstore";
 
     if (Record *VT = P.getMemoryVT())
       Explanation += (" MemVT=" + VT->getName()).str();
     if (Record *VT = P.getScalarMemoryVT())
       Explanation += (" ScalarVT(MemVT)=" + VT->getName()).str();
 
+    if (ListInit *AddrSpaces = P.getAddressSpaces()) {
+      bool First = true;
+      raw_string_ostream OS(Explanation);
+      OS << " AddressSpaces=[";
+
+      for (Init *Val : AddrSpaces->getValues()) {
+        IntInit *IntVal = dyn_cast<IntInit>(Val);
+        if (!IntVal)
+          continue;
+
+        if (First)
+          First = false;
+        else
+          OS << ',';
+
+        OS << IntVal->getValue();
+      }
+
+      OS << ']';
+    }
+
     if (P.isAtomicOrderingMonotonic())
       Explanation += " monotonic";
     if (P.isAtomicOrderingAcquire())
       Explanation += " acquire";
     if (P.isAtomicOrderingRelease())
       Explanation += " release";
     if (P.isAtomicOrderingAcquireRelease())
       Explanation += " acq_rel";
for (const TreePredicateCall &Call : N->getPredicateCalls()) {
     if (Predicate.isLoad() && Predicate.getMemoryVT())
       continue;
 
     if (Predicate.isLoad() || Predicate.isStore()) {
       if (Predicate.isUnindexed())
         continue;
     }
 
+    if (Predicate.isLoad() || Predicate.isStore() || Predicate.isAtomic()) {
+      const ListInit *AddrSpaces = Predicate.getAddressSpaces();
+      if (AddrSpaces && !AddrSpaces->empty())
+        continue;
+    }
+
     if (Predicate.isAtomic() && Predicate.getMemoryVT())
       continue;
 
     if (Predicate.isAtomic() &&
         (Predicate.isAtomicOrderingMonotonic() ||
          Predicate.isAtomicOrderingAcquire() ||
          Predicate.isAtomicOrderingRelease() ||
          Predicate.isAtomicOrderingAcquireRelease() ||
public:
   /// are currently not compared between each other.
   enum PredicateKind {
     IPM_Opcode,
     IPM_NumOperands,
     IPM_ImmPredicate,
     IPM_AtomicOrderingMMO,
     IPM_MemoryLLTSize,
     IPM_MemoryVsLLTSize,
+    IPM_MemoryAddressSpace,
     IPM_GenericPredicate,
     OPM_SameOperand,
     OPM_ComplexPattern,
     OPM_IntrinsicID,
     OPM_Instruction,
     OPM_Int,
     OPM_LiteralInt,
     OPM_LLT,
void emitPredicateOpcodes(MatchTable &Table,
     Table << MatchTable::Opcode("GIM_CheckMemorySizeEqualTo")
           << MatchTable::Comment("MI") << MatchTable::IntValue(InsnVarID)
           << MatchTable::Comment("MMO") << MatchTable::IntValue(MMOIdx)
           << MatchTable::Comment("Size") << MatchTable::IntValue(Size)
           << MatchTable::LineBreak;
   }
 };
 
+class MemoryAddressSpacePredicateMatcher : public InstructionPredicateMatcher {
+protected:
+  unsigned MMOIdx;
+  SmallVector<unsigned, 4> AddrSpaces;
+
+public:
+  MemoryAddressSpacePredicateMatcher(unsigned InsnVarID, unsigned MMOIdx,
+                                     ArrayRef<unsigned> AddrSpaces)
+      : InstructionPredicateMatcher(IPM_MemoryAddressSpace, InsnVarID),
+        MMOIdx(MMOIdx), AddrSpaces(AddrSpaces.begin(), AddrSpaces.end()) {}
+
+  static bool classof(const PredicateMatcher *P) {
+    return P->getKind() == IPM_MemoryAddressSpace;
+  }
+  bool isIdentical(const PredicateMatcher &B) const override {
+    if (!InstructionPredicateMatcher::isIdentical(B))
+      return false;
+    auto *Other = cast<MemoryAddressSpacePredicateMatcher>(&B);
+    return MMOIdx == Other->MMOIdx && AddrSpaces == Other->AddrSpaces;
+  }
+
+  void emitPredicateOpcodes(MatchTable &Table,
+                            RuleMatcher &Rule) const override {
+    Table << MatchTable::Opcode("GIM_CheckMemoryAddressSpace")
+          << MatchTable::Comment("MI") << MatchTable::IntValue(InsnVarID)
+          << MatchTable::Comment("MMO") << MatchTable::IntValue(MMOIdx)
+        // Encode number of address spaces to expect.
+          << MatchTable::Comment("NumAddrSpace")
+          << MatchTable::IntValue(AddrSpaces.size());
+    for (unsigned AS : AddrSpaces)
+      Table << MatchTable::Comment("AddrSpace") << MatchTable::IntValue(AS);
+
+    Table << MatchTable::LineBreak;
+  }
+};
+
 /// Generates code to check that the size of an MMO is less-than, equal-to, or
 /// greater than a given LLT.
 class MemoryVsLLTSizePredicateMatcher : public InstructionPredicateMatcher {
 public:
   enum RelationKind {
     GreaterThan,
     EqualTo,
     LessThan,
for (const TreePredicateCall &Call : Src->getPredicateCalls()) {
     if (Predicate.isAlwaysTrue())
       continue;
 
     if (Predicate.isImmediatePattern()) {
       InsnMatcher.addPredicate<InstructionImmPredicateMatcher>(Predicate);
       continue;
     }
 
-    // G_LOAD is used for both non-extending and any-extending loads. 
+    // An address space check is needed in all contexts if there is one.
+    if (Predicate.isLoad() || Predicate.isStore() || Predicate.isAtomic()) {
+      if (const ListInit *AddrSpaces = Predicate.getAddressSpaces()) {
+        SmallVector<unsigned, 4> ParsedAddrSpaces;
+
+        for (Init *Val : AddrSpaces->getValues()) {
+          IntInit *IntVal = dyn_cast<IntInit>(Val);
+          if (!IntVal)
+            return failedImport("Address space is not an integer");
+          ParsedAddrSpaces.push_back(IntVal->getValue());
+        }
+
+        if (!ParsedAddrSpaces.empty()) {
+          InsnMatcher.addPredicate<MemoryAddressSpacePredicateMatcher>(
+            0, ParsedAddrSpaces);
+        }
+      }
+    }
+
+    // G_LOAD is used for both non-extending and any-extending loads.
     if (Predicate.isLoad() && Predicate.isNonExtLoad()) {
       InsnMatcher.addPredicate<MemoryVsLLTSizePredicateMatcher>(
           0, MemoryVsLLTSizePredicateMatcher::EqualTo, 0);
       continue;
     }
     if (Predicate.isLoad() && Predicate.isAnyExtLoad()) {
       InsnMatcher.addPredicate<MemoryVsLLTSizePredicateMatcher>(
           0, MemoryVsLLTSizePredicateMatcher::LessThan, 0);