[globalisel][legalizerinfo] Add support for legalization based on the MachineMemOperand

Summary:
Currently only the memory size is supported but others can be added as
needed.

narrowScalar for G_LOAD and G_STORE now correctly update the
MachineMemOperand and will refuse to legalize atomics since those need more
careful expansions to maintain atomicity.

Reviewers: ab, aditya_nandakumar, bogner, rtereshin, aemerson, javed.absar

Reviewed By: aemerson

Subscribers: aemerson, rovka, kristof.beyls, javed.absar, llvm-commits

Differential Revision: https://reviews.llvm.org/D45466

llvm-svn: 331071
diff --git a/llvm/lib/CodeGen/GlobalISel/LegalityPredicates.cpp b/llvm/lib/CodeGen/GlobalISel/LegalityPredicates.cpp
index 0d30e07..d9ec88f 100644
--- a/llvm/lib/CodeGen/GlobalISel/LegalityPredicates.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/LegalityPredicates.cpp
@@ -41,6 +41,18 @@
   };
 }
 
+LegalityPredicate LegalityPredicates::typePairAndMemSizeInSet(
+    unsigned TypeIdx0, unsigned TypeIdx1, unsigned MMOIdx,
+    std::initializer_list<std::tuple<LLT, LLT, unsigned>> TypesAndMemSizeInit) {
+  SmallVector<std::tuple<LLT, LLT, unsigned>, 4> TypesAndMemSize = TypesAndMemSizeInit;
+  return [=](const LegalityQuery &Query) {
+    std::tuple<LLT, LLT, unsigned> Match = {
+        Query.Types[TypeIdx0], Query.Types[TypeIdx1], Query.MMODescrs[MMOIdx].Size};
+    return std::find(TypesAndMemSize.begin(), TypesAndMemSize.end(), Match) !=
+           TypesAndMemSize.end();
+  };
+}
+
 LegalityPredicate LegalityPredicates::isScalar(unsigned TypeIdx) {
   return [=](const LegalityQuery &Query) {
     return Query.Types[TypeIdx].isScalar();
@@ -70,6 +82,12 @@
   };
 }
 
+LegalityPredicate LegalityPredicates::memSizeInBytesNotPow2(unsigned MMOIdx) {
+  return [=](const LegalityQuery &Query) {
+    return !isPowerOf2_32(Query.MMODescrs[MMOIdx].Size /* In Bytes */);
+  };
+}
+
 LegalityPredicate LegalityPredicates::numElementsNotPow2(unsigned TypeIdx) {
   return [=](const LegalityQuery &Query) {
     const LLT &QueryTy = Query.Types[TypeIdx];
diff --git a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
index d8be22f..60901e0 100644
--- a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
@@ -272,8 +272,8 @@
 
   MIRBuilder.setInstr(MI);
 
-  int64_t SizeOp0 = MRI.getType(MI.getOperand(0).getReg()).getSizeInBits();
-  int64_t NarrowSize = NarrowTy.getSizeInBits();
+  uint64_t SizeOp0 = MRI.getType(MI.getOperand(0).getReg()).getSizeInBits();
+  uint64_t NarrowSize = NarrowTy.getSizeInBits();
 
   switch (MI.getOpcode()) {
   default:
@@ -339,8 +339,8 @@
     extractParts(MI.getOperand(1).getReg(), NarrowTy, NumParts, SrcRegs);
 
     unsigned OpReg = MI.getOperand(0).getReg();
-    int64_t OpStart = MI.getOperand(2).getImm();
-    int64_t OpSize = MRI.getType(OpReg).getSizeInBits();
+    uint64_t OpStart = MI.getOperand(2).getImm();
+    uint64_t OpSize = MRI.getType(OpReg).getSizeInBits();
     for (int i = 0; i < NumParts; ++i) {
       unsigned SrcStart = i * NarrowSize;
 
@@ -355,7 +355,8 @@
 
       // OpSegStart is where this destination segment would start in OpReg if it
       // extended infinitely in both directions.
-      int64_t ExtractOffset, SegSize;
+      int64_t ExtractOffset;
+      uint64_t SegSize;
       if (OpStart < SrcStart) {
         ExtractOffset = 0;
         SegSize = std::min(NarrowSize, OpStart + OpSize - SrcStart);
@@ -391,8 +392,8 @@
     extractParts(MI.getOperand(1).getReg(), NarrowTy, NumParts, SrcRegs);
 
     unsigned OpReg = MI.getOperand(2).getReg();
-    int64_t OpStart = MI.getOperand(3).getImm();
-    int64_t OpSize = MRI.getType(OpReg).getSizeInBits();
+    uint64_t OpStart = MI.getOperand(3).getImm();
+    uint64_t OpSize = MRI.getType(OpReg).getSizeInBits();
     for (int i = 0; i < NumParts; ++i) {
       unsigned DstStart = i * NarrowSize;
 
@@ -409,7 +410,8 @@
 
       // OpSegStart is where this destination segment would start in OpReg if it
       // extended infinitely in both directions.
-      int64_t ExtractOffset, InsertOffset, SegSize;
+      int64_t ExtractOffset, InsertOffset;
+      uint64_t SegSize;
       if (OpStart < DstStart) {
         InsertOffset = 0;
         ExtractOffset = DstStart - OpStart;
@@ -443,6 +445,14 @@
     // NarrowSize.
     if (SizeOp0 % NarrowSize != 0)
       return UnableToLegalize;
+
+    const auto &MMO = **MI.memoperands_begin();
+    // This implementation doesn't work for atomics. Give up instead of doing
+    // something invalid.
+    if (MMO.getOrdering() != AtomicOrdering::NotAtomic ||
+        MMO.getFailureOrdering() != AtomicOrdering::NotAtomic)
+      return UnableToLegalize;
+
     int NumParts = SizeOp0 / NarrowSize;
     LLT OffsetTy = LLT::scalar(
         MRI.getType(MI.getOperand(1).getReg()).getScalarSizeInBits());
@@ -453,12 +463,16 @@
       unsigned SrcReg = 0;
       unsigned Adjustment = i * NarrowSize / 8;
 
+      MachineMemOperand *SplitMMO = MIRBuilder.getMF().getMachineMemOperand(
+          MMO.getPointerInfo().getWithOffset(Adjustment), MMO.getFlags(),
+          NarrowSize / 8, i == 0 ? MMO.getAlignment() : NarrowSize / 8,
+          MMO.getAAInfo(), MMO.getRanges(), MMO.getSyncScopeID(),
+          MMO.getOrdering(), MMO.getFailureOrdering());
+
       MIRBuilder.materializeGEP(SrcReg, MI.getOperand(1).getReg(), OffsetTy,
                                 Adjustment);
 
-      // TODO: This is conservatively correct, but we probably want to split the
-      // memory operands in the future.
-      MIRBuilder.buildLoad(DstReg, SrcReg, **MI.memoperands_begin());
+      MIRBuilder.buildLoad(DstReg, SrcReg, *SplitMMO);
 
       DstRegs.push_back(DstReg);
     }
@@ -472,6 +486,14 @@
     // NarrowSize.
     if (SizeOp0 % NarrowSize != 0)
       return UnableToLegalize;
+
+    const auto &MMO = **MI.memoperands_begin();
+    // This implementation doesn't work for atomics. Give up instead of doing
+    // something invalid.
+    if (MMO.getOrdering() != AtomicOrdering::NotAtomic ||
+        MMO.getFailureOrdering() != AtomicOrdering::NotAtomic)
+      return UnableToLegalize;
+
     int NumParts = SizeOp0 / NarrowSize;
     LLT OffsetTy = LLT::scalar(
         MRI.getType(MI.getOperand(1).getReg()).getScalarSizeInBits());
@@ -483,12 +505,16 @@
       unsigned DstReg = 0;
       unsigned Adjustment = i * NarrowSize / 8;
 
+      MachineMemOperand *SplitMMO = MIRBuilder.getMF().getMachineMemOperand(
+          MMO.getPointerInfo().getWithOffset(Adjustment), MMO.getFlags(),
+          NarrowSize / 8, i == 0 ? MMO.getAlignment() : NarrowSize / 8,
+          MMO.getAAInfo(), MMO.getRanges(), MMO.getSyncScopeID(),
+          MMO.getOrdering(), MMO.getFailureOrdering());
+
       MIRBuilder.materializeGEP(DstReg, MI.getOperand(1).getReg(), OffsetTy,
                                 Adjustment);
 
-      // TODO: This is conservatively correct, but we probably want to split the
-      // memory operands in the future.
-      MIRBuilder.buildStore(SrcRegs[i], DstReg, **MI.memoperands_begin());
+      MIRBuilder.buildStore(SrcRegs[i], DstReg, *SplitMMO);
     }
     MI.eraseFromParent();
     return Legalized;
diff --git a/llvm/lib/CodeGen/GlobalISel/LegalizerInfo.cpp b/llvm/lib/CodeGen/GlobalISel/LegalizerInfo.cpp
index 32fec80..a681e9f 100644
--- a/llvm/lib/CodeGen/GlobalISel/LegalizerInfo.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/LegalizerInfo.cpp
@@ -42,11 +42,18 @@
     cl::Hidden);
 
 raw_ostream &LegalityQuery::print(raw_ostream &OS) const {
-  OS << Opcode << ", {";
+  OS << Opcode << ", Tys={";
   for (const auto &Type : Types) {
     OS << Type << ", ";
   }
+  OS << "}, Opcode=";
+
+  OS << Opcode << ", MMOs={";
+  for (const auto &MMODescr : MMODescrs) {
+    OS << MMODescr.Size << ", ";
+  }
   OS << "}";
+
   return OS;
 }
 
@@ -330,7 +337,12 @@
     LLT Ty = getTypeFromTypeIdx(MI, MRI, i, TypeIdx);
     Types.push_back(Ty);
   }
-  return getAction({MI.getOpcode(), Types});
+
+  SmallVector<LegalityQuery::MemDesc, 2> MemDescrs;
+  for (const auto &MMO : MI.memoperands())
+    MemDescrs.push_back({MMO->getSize() /* in bytes */ * 8});
+
+  return getAction({MI.getOpcode(), Types, MemDescrs});
 }
 
 bool LegalizerInfo::isLegal(const MachineInstr &MI,