Reuse-EmitZAInstr-to-add-Za-Matrix

CarolineConcatto · CarolineConcatto · commit f0e5e45ec660 · 2024-06-04T14:14:58.000Z
diff --git a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
@@ -395,8 +395,8 @@ class AArch64DAGToDAGISel : public SelectionDAGISel {
   template <unsigned MaxIdx, unsigned Scale>
   void SelectMultiVectorMove(SDNode *N, unsigned NumVecs, unsigned BaseReg,
                              unsigned Op);
-  template <unsigned MaxIdx, unsigned Scale>
-  void SelectMultiVectorMoveZ(SDNode *N, unsigned NumVecs, unsigned Op);
+  void SelectMultiVectorMoveZ(SDNode *N, unsigned NumVecs, unsigned Op,
+                              unsigned MaxIdx, unsigned Scale);
   bool SelectAddrModeFrameIndexSVE(SDValue N, SDValue &Base, SDValue &OffImm);
   /// SVE Reg+Imm addressing mode.
   template <int64_t Min, int64_t Max>
@@ -2004,9 +2004,9 @@ void AArch64DAGToDAGISel::SelectMultiVectorMove(SDNode *N, unsigned NumVecs,
   CurDAG->RemoveDeadNode(N);
 }
 
-template <unsigned MaxIdx, unsigned Scale>
 void AArch64DAGToDAGISel::SelectMultiVectorMoveZ(SDNode *N, unsigned NumVecs,
-                                                 unsigned Op) {
+                                                 unsigned Op, unsigned MaxIdx,
+                                                 unsigned Scale) {
 
   SDValue SliceBase = N->getOperand(3);
   SDValue Base, Offset;
@@ -5274,68 +5274,68 @@ void AArch64DAGToDAGISel::Select(SDNode *Node) {
     }
     case Intrinsic::aarch64_sme_readz_horiz_x2: {
       if (VT == MVT::nxv16i8) {
-        SelectMultiVectorMoveZ<14, 2>(Node, 2, AArch64::MOVAZ_2ZMI_H_B_PSEUDO);
+        SelectMultiVectorMoveZ(Node, 2, AArch64::MOVAZ_2ZMI_H_B_PSEUDO, 14, 2);
         return;
       } else if (VT == MVT::nxv8i16 || VT == MVT::nxv8f16 ||
                  VT == MVT::nxv8bf16) {
-        SelectMultiVectorMoveZ<6, 2>(Node, 2, AArch64::MOVAZ_2ZMI_H_H_PSEUDO);
+        SelectMultiVectorMoveZ(Node, 2, AArch64::MOVAZ_2ZMI_H_H_PSEUDO, 6, 2);
         return;
       } else if (VT == MVT::nxv4i32 || VT == MVT::nxv4f32) {
-        SelectMultiVectorMoveZ<2, 2>(Node, 2, AArch64::MOVAZ_2ZMI_H_S_PSEUDO);
+        SelectMultiVectorMoveZ(Node, 2, AArch64::MOVAZ_2ZMI_H_S_PSEUDO, 2, 2);
         return;
       } else if (VT == MVT::nxv2i64 || VT == MVT::nxv2f64) {
-        SelectMultiVectorMoveZ<0, 2>(Node, 2, AArch64::MOVAZ_2ZMI_H_D_PSEUDO);
+        SelectMultiVectorMoveZ(Node, 2, AArch64::MOVAZ_2ZMI_H_D_PSEUDO, 0, 2);
         return;
       }
       break;
     }
     case Intrinsic::aarch64_sme_readz_vert_x2: {
       if (VT == MVT::nxv16i8) {
-        SelectMultiVectorMoveZ<14, 2>(Node, 2, AArch64::MOVAZ_2ZMI_V_B_PSEUDO);
+        SelectMultiVectorMoveZ(Node, 2, AArch64::MOVAZ_2ZMI_V_B_PSEUDO, 14, 2);
         return;
       } else if (VT == MVT::nxv8i16 || VT == MVT::nxv8f16 ||
                  VT == MVT::nxv8bf16) {
-        SelectMultiVectorMoveZ<6, 2>(Node, 2, AArch64::MOVAZ_2ZMI_V_H_PSEUDO);
+        SelectMultiVectorMoveZ(Node, 2, AArch64::MOVAZ_2ZMI_V_H_PSEUDO, 6, 2);
         return;
       } else if (VT == MVT::nxv4i32 || VT == MVT::nxv4f32) {
-        SelectMultiVectorMoveZ<2, 2>(Node, 2, AArch64::MOVAZ_2ZMI_V_S_PSEUDO);
+        SelectMultiVectorMoveZ(Node, 2, AArch64::MOVAZ_2ZMI_V_S_PSEUDO, 2, 2);
         return;
       } else if (VT == MVT::nxv2i64 || VT == MVT::nxv2f64) {
-        SelectMultiVectorMoveZ<0, 2>(Node, 2, AArch64::MOVAZ_2ZMI_V_D_PSEUDO);
+        SelectMultiVectorMoveZ(Node, 2, AArch64::MOVAZ_2ZMI_V_D_PSEUDO, 0, 2);
         return;
       }
       break;
     }
     case Intrinsic::aarch64_sme_readz_horiz_x4: {
       if (VT == MVT::nxv16i8) {
-        SelectMultiVectorMoveZ<12, 4>(Node, 4, AArch64::MOVAZ_4ZMI_H_B_PSEUDO);
+        SelectMultiVectorMoveZ(Node, 4, AArch64::MOVAZ_4ZMI_H_B_PSEUDO, 12, 4);
         return;
       } else if (VT == MVT::nxv8i16 || VT == MVT::nxv8f16 ||
                  VT == MVT::nxv8bf16) {
-        SelectMultiVectorMoveZ<4, 4>(Node, 4, AArch64::MOVAZ_4ZMI_H_H_PSEUDO);
+        SelectMultiVectorMoveZ(Node, 4, AArch64::MOVAZ_4ZMI_H_H_PSEUDO, 4, 4);
         return;
       } else if (VT == MVT::nxv4i32 || VT == MVT::nxv4f32) {
-        SelectMultiVectorMoveZ<0, 4>(Node, 4, AArch64::MOVAZ_4ZMI_H_S_PSEUDO);
+        SelectMultiVectorMoveZ(Node, 4, AArch64::MOVAZ_4ZMI_H_S_PSEUDO, 0, 4);
         return;
       } else if (VT == MVT::nxv2i64 || VT == MVT::nxv2f64) {
-        SelectMultiVectorMoveZ<0, 4>(Node, 4, AArch64::MOVAZ_4ZMI_H_D_PSEUDO);
+        SelectMultiVectorMoveZ(Node, 4, AArch64::MOVAZ_4ZMI_H_D_PSEUDO, 0, 4);
         return;
       }
       break;
     }
     case Intrinsic::aarch64_sme_readz_vert_x4: {
       if (VT == MVT::nxv16i8) {
-        SelectMultiVectorMoveZ<12, 4>(Node, 4, AArch64::MOVAZ_4ZMI_V_B_PSEUDO);
+        SelectMultiVectorMoveZ(Node, 4, AArch64::MOVAZ_4ZMI_V_B_PSEUDO, 12, 4);
         return;
       } else if (VT == MVT::nxv8i16 || VT == MVT::nxv8f16 ||
                  VT == MVT::nxv8bf16) {
-        SelectMultiVectorMoveZ<4, 4>(Node, 4, AArch64::MOVAZ_4ZMI_V_H_PSEUDO);
+        SelectMultiVectorMoveZ(Node, 4, AArch64::MOVAZ_4ZMI_V_H_PSEUDO, 4, 4);
         return;
       } else if (VT == MVT::nxv4i32 || VT == MVT::nxv4f32) {
-        SelectMultiVectorMoveZ<0, 4>(Node, 4, AArch64::MOVAZ_4ZMI_V_S_PSEUDO);
+        SelectMultiVectorMoveZ(Node, 4, AArch64::MOVAZ_4ZMI_V_S_PSEUDO, 0, 4);
         return;
       } else if (VT == MVT::nxv2i64 || VT == MVT::nxv2f64) {
-        SelectMultiVectorMoveZ<0, 4>(Node, 4, AArch64::MOVAZ_4ZMI_V_D_PSEUDO);
+        SelectMultiVectorMoveZ(Node, 4, AArch64::MOVAZ_4ZMI_V_D_PSEUDO, 0, 4);
         return;
       }
       break;
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -2965,19 +2965,28 @@ MachineBasicBlock *AArch64TargetLowering::EmitZTInstr(MachineInstr &MI,
 
 MachineBasicBlock *
 AArch64TargetLowering::EmitZAInstr(unsigned Opc, unsigned BaseReg,
-                                   MachineInstr &MI,
-                                   MachineBasicBlock *BB, bool HasTile) const {
+                                   MachineInstr &MI, MachineBasicBlock *BB,
+                                   bool HasTile, bool HasZPROut) const {
   const TargetInstrInfo *TII = Subtarget->getInstrInfo();
   MachineInstrBuilder MIB = BuildMI(*BB, MI, MI.getDebugLoc(), TII->get(Opc));
   unsigned StartIdx = 0;
 
-  if (HasTile) {
-    MIB.addReg(BaseReg + MI.getOperand(0).getImm(), RegState::Define);
-    MIB.addReg(BaseReg + MI.getOperand(0).getImm());
-    StartIdx = 1;
-  } else
-    MIB.addReg(BaseReg, RegState::Define).addReg(BaseReg);
-
+  if (HasZPROut) {
+    if (HasTile) {
+      MIB.add(MI.getOperand(0)); // Output ZPR
+      MIB.addReg(BaseReg + MI.getOperand(1).getImm(),
+                 RegState::Define);                    // Output ZA Tile
+      MIB.addReg(BaseReg + MI.getOperand(1).getImm()); // Input Za Tile
+      StartIdx = 2;
+    }
+  } else {
+    if (HasTile) {
+      MIB.addReg(BaseReg + MI.getOperand(0).getImm(), RegState::Define);
+      MIB.addReg(BaseReg + MI.getOperand(0).getImm());
+      StartIdx = 1;
+    } else
+      MIB.addReg(BaseReg, RegState::Define).addReg(BaseReg);
+  }
   for (unsigned I = StartIdx; I < MI.getNumOperands(); ++I)
     MIB.add(MI.getOperand(I));
 
@@ -3012,17 +3021,59 @@ MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
         TII->get(MI.getOpcode()).TSFlags & AArch64::SMEMatrixTypeMask;
     switch (SMEMatrixType) {
     case (AArch64::SMEMatrixArray):
-      return EmitZAInstr(SMEOrigInstr, AArch64::ZA, MI, BB, /*HasTile*/ false);
+      return EmitZAInstr(SMEOrigInstr, AArch64::ZA, MI, BB, /*HasTile*/ false,
+                         /*HasZPROut*/ false);
     case (AArch64::SMEMatrixTileB):
-      return EmitZAInstr(SMEOrigInstr, AArch64::ZAB0, MI, BB, /*HasTile*/ true);
+      switch (MI.getOpcode()) {
+      case AArch64::MOVAZ_2ZMI_H_B_PSEUDO:
+      case AArch64::MOVAZ_2ZMI_V_B_PSEUDO:
+      case AArch64::MOVAZ_4ZMI_H_B_PSEUDO:
+      case AArch64::MOVAZ_4ZMI_V_B_PSEUDO:
+        return EmitZAInstr(SMEOrigInstr, AArch64::ZAB0, MI, BB,
+                           /*HasTile*/ true, /*HasZPROut*/ true);
+      default:
+        return EmitZAInstr(SMEOrigInstr, AArch64::ZAB0, MI, BB,
+                           /*HasTile*/ true, /*HasZPROut*/ false);
+      }
     case (AArch64::SMEMatrixTileH):
-      return EmitZAInstr(SMEOrigInstr, AArch64::ZAH0, MI, BB, /*HasTile*/ true);
+      switch (MI.getOpcode()) {
+      case AArch64::MOVAZ_2ZMI_H_H_PSEUDO:
+      case AArch64::MOVAZ_2ZMI_V_H_PSEUDO:
+      case AArch64::MOVAZ_4ZMI_H_H_PSEUDO:
+      case AArch64::MOVAZ_4ZMI_V_H_PSEUDO:
+        return EmitZAInstr(SMEOrigInstr, AArch64::ZAH0, MI, BB,
+                           /*HasTile*/ true, /*HasZPROut*/ true);
+      default:
+        return EmitZAInstr(SMEOrigInstr, AArch64::ZAH0, MI, BB,
+                           /*HasTile*/ true, /*HasZPROut*/ false);
+      }
     case (AArch64::SMEMatrixTileS):
-      return EmitZAInstr(SMEOrigInstr, AArch64::ZAS0, MI, BB, /*HasTile*/ true);
+      switch (MI.getOpcode()) {
+      case AArch64::MOVAZ_2ZMI_H_S_PSEUDO:
+      case AArch64::MOVAZ_2ZMI_V_S_PSEUDO:
+      case AArch64::MOVAZ_4ZMI_H_S_PSEUDO:
+      case AArch64::MOVAZ_4ZMI_V_S_PSEUDO:
+        return EmitZAInstr(SMEOrigInstr, AArch64::ZAS0, MI, BB,
+                           /*HasTile*/ true, /*HasZPROut*/ true);
+      default:
+        return EmitZAInstr(SMEOrigInstr, AArch64::ZAS0, MI, BB,
+                           /*HasTile*/ true, /*HasZPROut*/ false);
+      }
     case (AArch64::SMEMatrixTileD):
-      return EmitZAInstr(SMEOrigInstr, AArch64::ZAD0, MI, BB, /*HasTile*/ true);
+      switch (MI.getOpcode()) {
+      case AArch64::MOVAZ_2ZMI_H_D_PSEUDO:
+      case AArch64::MOVAZ_2ZMI_V_D_PSEUDO:
+      case AArch64::MOVAZ_4ZMI_H_D_PSEUDO:
+      case AArch64::MOVAZ_4ZMI_V_D_PSEUDO:
+        return EmitZAInstr(SMEOrigInstr, AArch64::ZAD0, MI, BB,
+                           /*HasTile*/ true, /*HasZPROut*/ true);
+      default:
+        return EmitZAInstr(SMEOrigInstr, AArch64::ZAD0, MI, BB,
+                           /*HasTile*/ true, /*HasZPROut*/ false);
+      }
     case (AArch64::SMEMatrixTileQ):
-      return EmitZAInstr(SMEOrigInstr, AArch64::ZAQ0, MI, BB, /*HasTile*/ true);
+      return EmitZAInstr(SMEOrigInstr, AArch64::ZAQ0, MI, BB, /*HasTile*/ true,
+                         /*HasZPROut*/ false);
     }
   }
 
@@ -3091,38 +3142,6 @@ MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
     return EmitZero(MI, BB);
   case AArch64::ZERO_T_PSEUDO:
     return EmitZTInstr(MI, BB, AArch64::ZERO_T, /*Op0IsDef=*/true);
-  case AArch64::MOVAZ_2ZMI_H_B_PSEUDO:
-    return EmitTileMovaz(AArch64::MOVAZ_2ZMI_H_B, AArch64::ZAB0, MI, BB);
-  case AArch64::MOVAZ_2ZMI_H_H_PSEUDO:
-    return EmitTileMovaz(AArch64::MOVAZ_2ZMI_H_H, AArch64::ZAH0, MI, BB);
-  case AArch64::MOVAZ_2ZMI_H_S_PSEUDO:
-    return EmitTileMovaz(AArch64::MOVAZ_2ZMI_H_S, AArch64::ZAS0, MI, BB);
-  case AArch64::MOVAZ_2ZMI_H_D_PSEUDO:
-    return EmitTileMovaz(AArch64::MOVAZ_2ZMI_H_D, AArch64::ZAD0, MI, BB);
-  case AArch64::MOVAZ_2ZMI_V_B_PSEUDO:
-    return EmitTileMovaz(AArch64::MOVAZ_2ZMI_V_B, AArch64::ZAB0, MI, BB);
-  case AArch64::MOVAZ_2ZMI_V_H_PSEUDO:
-    return EmitTileMovaz(AArch64::MOVAZ_2ZMI_V_H, AArch64::ZAH0, MI, BB);
-  case AArch64::MOVAZ_2ZMI_V_S_PSEUDO:
-    return EmitTileMovaz(AArch64::MOVAZ_2ZMI_V_S, AArch64::ZAS0, MI, BB);
-  case AArch64::MOVAZ_2ZMI_V_D_PSEUDO:
-    return EmitTileMovaz(AArch64::MOVAZ_2ZMI_V_D, AArch64::ZAD0, MI, BB);
-  case AArch64::MOVAZ_4ZMI_H_B_PSEUDO:
-    return EmitTileMovaz(AArch64::MOVAZ_4ZMI_H_B, AArch64::ZAB0, MI, BB);
-  case AArch64::MOVAZ_4ZMI_H_H_PSEUDO:
-    return EmitTileMovaz(AArch64::MOVAZ_4ZMI_H_H, AArch64::ZAH0, MI, BB);
-  case AArch64::MOVAZ_4ZMI_H_S_PSEUDO:
-    return EmitTileMovaz(AArch64::MOVAZ_4ZMI_H_S, AArch64::ZAS0, MI, BB);
-  case AArch64::MOVAZ_4ZMI_H_D_PSEUDO:
-    return EmitTileMovaz(AArch64::MOVAZ_4ZMI_H_D, AArch64::ZAD0, MI, BB);
-  case AArch64::MOVAZ_4ZMI_V_B_PSEUDO:
-    return EmitTileMovaz(AArch64::MOVAZ_4ZMI_V_B, AArch64::ZAB0, MI, BB);
-  case AArch64::MOVAZ_4ZMI_V_H_PSEUDO:
-    return EmitTileMovaz(AArch64::MOVAZ_4ZMI_V_H, AArch64::ZAH0, MI, BB);
-  case AArch64::MOVAZ_4ZMI_V_S_PSEUDO:
-    return EmitTileMovaz(AArch64::MOVAZ_4ZMI_V_S, AArch64::ZAS0, MI, BB);
-  case AArch64::MOVAZ_4ZMI_V_D_PSEUDO:
-    return EmitTileMovaz(AArch64::MOVAZ_4ZMI_V_D, AArch64::ZAD0, MI, BB);
   }
 }
 
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.h b/llvm/lib/Target/AArch64/AArch64ISelLowering.h
@@ -654,7 +654,7 @@ class AArch64TargetLowering : public TargetLowering {
   MachineBasicBlock *EmitFill(MachineInstr &MI, MachineBasicBlock *BB) const;
   MachineBasicBlock *EmitZAInstr(unsigned Opc, unsigned BaseReg,
                                  MachineInstr &MI, MachineBasicBlock *BB,
-                                 bool HasTile) const;
+                                 bool HasTile, bool HasZPROut) const;
   MachineBasicBlock *EmitZTInstr(MachineInstr &MI, MachineBasicBlock *BB,
                                  unsigned Opcode, bool Op0IsDef) const;
   MachineBasicBlock *EmitZero(MachineInstr &MI, MachineBasicBlock *BB) const;
diff --git a/llvm/lib/Target/AArch64/SMEInstrFormats.td b/llvm/lib/Target/AArch64/SMEInstrFormats.td
