[CodeGen] Add MO_LaneMask type and a new COPY_LANEMASK instruction (#151944)

Introduce MO_LaneMask as new machine operand type. This can be used to
hold liveness infomation at sub-register granularity for register-type
operands. We also introduce a new COPY_LANEMASK instruction that uses
MO_lanemask operand to perform partial copy from source register
opernad.

One such use case of MO_LaneMask can be seen in #151123, where it can be
used to store live regUnits information corresponding to the source
register of the COPY instructions, later can be used during CopyPhysReg
expansion.

Author: vg0204

llvm/docs/MIRLangRef.rst

@@ -807,6 +807,22 @@ For an int eq predicate ``ICMP_EQ``, the syntax is:
 
    %2:gpr(s32) = G_ICMP intpred(eq), %0, %1
 
+LaneMask Operands
+^^^^^^^^^^^^^^^^^
+
+A LaneMask operand contains a LaneBitmask struct representing the covering of a
+register with sub-registers. Instructions typically associate a LaneMask operand
+with one or more Register operands, and use it to represent sub-register
+granularity information like liveness for those associated Register operands.
+
+
+For example, the COPY_LANEMASK instruction uses this operand to copy only active
+lanes (of the source register) in the mask. The syntax for it would look like:
+
+.. code-block:: text
+
+   $vgpr1 = COPY_LANEMASK $vgpr0, lanemask(0x00000000000000C0)
+
 .. TODO: Describe the parsers default behaviour when optional YAML attributes
    are missing.
 .. TODO: Describe the syntax for virtual register YAML definitions.

llvm/include/llvm/CodeGen/MIR2Vec.h

@@ -139,7 +139,7 @@ class MIRVocabulary {
       "FrameIndex",      "ConstantPoolIndex", "TargetIndex",  "JumpTableIndex",
       "ExternalSymbol",  "GlobalAddress",     "BlockAddress", "RegisterMask",
       "RegisterLiveOut", "Metadata",          "MCSymbol",     "CFIIndex",
-      "IntrinsicID",     "Predicate",         "ShuffleMask"};
+      "IntrinsicID",     "Predicate",         "ShuffleMask",  "LaneMask"};
   static_assert(std::size(CommonOperandNames) == MachineOperand::MO_Last - 1 &&
                 "Common operand names size changed, update accordingly");
 

llvm/include/llvm/CodeGen/MachineInstr.h

@@ -1451,6 +1451,10 @@ class MachineInstr
     return getOpcode() == TargetOpcode::COPY;
   }
 
+  bool isCopyLaneMask() const {
+    return getOpcode() == TargetOpcode::COPY_LANEMASK;
+  }
+
   bool isFullCopy() const {
     return isCopy() && !getOperand(0).getSubReg() && !getOperand(1).getSubReg();
   }
@@ -1484,6 +1488,7 @@ class MachineInstr
     case TargetOpcode::PHI:
     case TargetOpcode::G_PHI:
     case TargetOpcode::COPY:
+    case TargetOpcode::COPY_LANEMASK:
     case TargetOpcode::INSERT_SUBREG:
     case TargetOpcode::SUBREG_TO_REG:
     case TargetOpcode::REG_SEQUENCE:

llvm/include/llvm/CodeGen/MachineInstrBuilder.h

@@ -307,6 +307,11 @@ class MachineInstrBuilder {
     return *this;
   }
 
+  const MachineInstrBuilder &addLaneMask(LaneBitmask LaneMask) const {
+    MI->addOperand(*MF, MachineOperand::CreateLaneMask(LaneMask));
+    return *this;
+  }
+
   const MachineInstrBuilder &addSym(MCSymbol *Sym,
                                     unsigned char TargetFlags = 0) const {
     MI->addOperand(*MF, MachineOperand::CreateMCSymbol(Sym, TargetFlags));

llvm/include/llvm/CodeGen/MachineOperand.h

@@ -16,6 +16,7 @@
 #include "llvm/ADT/DenseMapInfo.h"
 #include "llvm/CodeGen/Register.h"
 #include "llvm/IR/Intrinsics.h"
+#include "llvm/MC/LaneBitmask.h"
 #include "llvm/Support/Compiler.h"
 #include <cassert>
 
@@ -69,7 +70,8 @@ class MachineOperand {
     MO_Predicate,         ///< Generic predicate for ISel
     MO_ShuffleMask,       ///< Other IR Constant for ISel (shuffle masks)
     MO_DbgInstrRef, ///< Integer indices referring to an instruction+operand
-    MO_Last = MO_DbgInstrRef
+    MO_LaneMask,    ///< Mask to represent active parts of registers
+    MO_Last = MO_LaneMask
   };
 
 private:
@@ -178,6 +180,7 @@ class MachineOperand {
     Intrinsic::ID IntrinsicID; // For MO_IntrinsicID.
     unsigned Pred;           // For MO_Predicate
     ArrayRef<int> ShuffleMask; // For MO_ShuffleMask
+    LaneBitmask LaneMask;      // For MO_LaneMask
 
     struct {                  // For MO_Register.
       // Register number is in SmallContents.RegNo.
@@ -360,6 +363,7 @@ class MachineOperand {
   bool isIntrinsicID() const { return OpKind == MO_IntrinsicID; }
   bool isPredicate() const { return OpKind == MO_Predicate; }
   bool isShuffleMask() const { return OpKind == MO_ShuffleMask; }
+  bool isLaneMask() const { return OpKind == MO_LaneMask; }
   //===--------------------------------------------------------------------===//
   // Accessors for Register Operands
   //===--------------------------------------------------------------------===//
@@ -624,6 +628,11 @@ class MachineOperand {
     return Contents.ShuffleMask;
   }
 
+  LaneBitmask getLaneMask() const {
+    assert(isLaneMask() && "Wrong MachineOperand accessor");
+    return Contents.LaneMask;
+  }
+
   /// Return the offset from the symbol in this operand. This always returns 0
   /// for ExternalSymbol operands.
   int64_t getOffset() const {
@@ -992,6 +1001,12 @@ class MachineOperand {
     return Op;
   }
 
+  static MachineOperand CreateLaneMask(LaneBitmask LaneMask) {
+    MachineOperand Op(MachineOperand::MO_LaneMask);
+    Op.Contents.LaneMask = LaneMask;
+    return Op;
+  }
+
   friend class MachineInstr;
   friend class MachineRegisterInfo;
 

llvm/include/llvm/Support/TargetOpcodes.def

@@ -114,6 +114,12 @@ HANDLE_TARGET_OPCODE(REG_SEQUENCE)
 /// used to copy between subregisters of virtual registers.
 HANDLE_TARGET_OPCODE(COPY)
 
+/// COPY_LANEMASK - Target-independent partial register copy. The laneMask
+/// operand indicates which parts of the source register are copied to the
+/// destination. Other parts of the destination are undefined. It does not
+/// support copy between virtual registers having subregister indices.
+HANDLE_TARGET_OPCODE(COPY_LANEMASK)
+
 /// BUNDLE - This instruction represents an instruction bundle. Instructions
 /// which immediately follow a BUNDLE instruction which are marked with
 /// 'InsideBundle' flag are inside the bundle.

llvm/include/llvm/Target/Target.td

@@ -1352,6 +1352,13 @@ def COPY : StandardPseudoInstruction {
   let isAsCheapAsAMove = true;
   let hasNoSchedulingInfo = false;
 }
+def COPY_LANEMASK : StandardPseudoInstruction {
+  let OutOperandList = (outs unknown:$dst);
+  let InOperandList = (ins unknown:$src, unknown:$lanemask);
+  let AsmString = "";
+  let hasSideEffects = false;
+  let isAsCheapAsAMove = true;
+}
 def BUNDLE : StandardPseudoInstruction {
   let OutOperandList = (outs);
   let InOperandList = (ins variable_ops);

llvm/lib/CodeGen/MIRParser/MILexer.cpp

@@ -266,6 +266,7 @@ static MIToken::TokenKind getIdentifierKind(StringRef Identifier) {
       .Case("constant-pool", MIToken::kw_constant_pool)
       .Case("call-entry", MIToken::kw_call_entry)
       .Case("custom", MIToken::kw_custom)
+      .Case("lanemask", MIToken::kw_lanemask)
       .Case("liveout", MIToken::kw_liveout)
       .Case("landing-pad", MIToken::kw_landing_pad)
       .Case("inlineasm-br-indirect-target",

llvm/lib/CodeGen/MIRParser/MILexer.h

@@ -122,6 +122,7 @@ struct MIToken {
     kw_constant_pool,
     kw_call_entry,
     kw_custom,
+    kw_lanemask,
     kw_liveout,
     kw_landing_pad,
     kw_inlineasm_br_indirect_target,

llvm/lib/CodeGen/MIRParser/MIParser.cpp

@@ -496,6 +496,7 @@ class MIParser {
   bool parseTargetIndexOperand(MachineOperand &Dest);
   bool parseDbgInstrRefOperand(MachineOperand &Dest);
   bool parseCustomRegisterMaskOperand(MachineOperand &Dest);
+  bool parseLaneMaskOperand(MachineOperand &Dest);
   bool parseLiveoutRegisterMaskOperand(MachineOperand &Dest);
   bool parseMachineOperand(const unsigned OpCode, const unsigned OpIdx,
                            MachineOperand &Dest,
@@ -2886,6 +2887,31 @@ bool MIParser::parseCustomRegisterMaskOperand(MachineOperand &Dest) {
   return false;
 }
 
+bool MIParser::parseLaneMaskOperand(MachineOperand &Dest) {
+  assert(Token.is(MIToken::kw_lanemask));
+
+  lex();
+  if (expectAndConsume(MIToken::lparen))
+    return true;
+
+  // Parse lanemask.
+  if (Token.isNot(MIToken::IntegerLiteral) && Token.isNot(MIToken::HexLiteral))
+    return error("expected a valid lane mask value");
+  static_assert(sizeof(LaneBitmask::Type) == sizeof(uint64_t),
+                "Use correct get-function for lane mask.");
+  LaneBitmask::Type V;
+  if (getUint64(V))
+    return true;
+  LaneBitmask LaneMask(V);
+  lex();
+
+  if (expectAndConsume(MIToken::rparen))
+    return true;
+
+  Dest = MachineOperand::CreateLaneMask(LaneMask);
+  return false;
+}
+
 bool MIParser::parseLiveoutRegisterMaskOperand(MachineOperand &Dest) {
   assert(Token.is(MIToken::kw_liveout));
   uint32_t *Mask = MF.allocateRegMask();
@@ -2986,6 +3012,8 @@ bool MIParser::parseMachineOperand(const unsigned OpCode, const unsigned OpIdx,
     return parseIntrinsicOperand(Dest);
   case MIToken::kw_target_index:
     return parseTargetIndexOperand(Dest);
+  case MIToken::kw_lanemask:
+    return parseLaneMaskOperand(Dest);
   case MIToken::kw_liveout:
     return parseLiveoutRegisterMaskOperand(Dest);
   case MIToken::kw_floatpred: