[AIE2] New vector custom post combiners

Optimize some vector operations (instrinsics) to COPYs

Author: andcarminati

llvm/lib/Target/AIE/AIE2InstrInfo.cpp

@@ -1321,3 +1321,87 @@ bool AIE2InstrInfo::canHoistCheapInst(const MachineInstr &MI) const {
          AIE2::eSRegClass.hasSubClassEq(
              MRI.getRegClass(MI.getOperand(0).getReg()));
 }
+
+std::optional<const AIEBaseInstrInfo::VConcatOpInfo>
+AIE2InstrInfo::getVConcatOpInfo(const MachineInstr &MI) const {
+
+  std::optional<const AIEBaseInstrInfo::VConcatOpInfo> BaseOpInfo =
+      AIEBaseInstrInfo::getVConcatOpInfo(MI);
+
+  if (BaseOpInfo)
+    return BaseOpInfo;
+
+  if (MI.getOpcode() != TargetOpcode::G_INTRINSIC)
+    return std::nullopt;
+
+  const GIntrinsic &GMI = cast<const GIntrinsic>(MI);
+
+  switch (GMI.getIntrinsicID()) {
+  case Intrinsic::aie2_concat_I512_I256:
+  case Intrinsic::aie2_concat_I1024_I512:
+  case Intrinsic::aie2_concat_I1024_I256:
+
+  case Intrinsic::aie2_concat_bf512_bf256:
+  case Intrinsic::aie2_concat_bf1024_bf512:
+  case Intrinsic::aie2_concat_bf1024_bf256:
+
+  case Intrinsic::aie2_concat_512_256_acc:
+  case Intrinsic::aie2_concat_1024_512_acc:
+  case Intrinsic::aie2_concat_1024_256_acc:
+    return VConcatOpInfo{2, 1};
+  default:
+    return std::nullopt;
+  }
+}
+
+std::optional<const AIEBaseInstrInfo::VUpdateOpInfo>
+AIE2InstrInfo::getVUpdateOpInfo(const MachineInstr &MI) const {
+
+  if (MI.getOpcode() != TargetOpcode::G_INTRINSIC)
+    return std::nullopt;
+
+  const GIntrinsic &GMI = cast<const GIntrinsic>(MI);
+
+  switch (GMI.getIntrinsicID()) {
+  case Intrinsic::aie2_upd_I512_I256:
+  case Intrinsic::aie2_upd_I1024_I512:
+  case Intrinsic::aie2_upd_I1024_I256:
+
+  case Intrinsic::aie2_upd_bf512_bf256:
+  case Intrinsic::aie2_upd_bf1024_bf512:
+  case Intrinsic::aie2_upd_bf1024_bf256:
+
+  case Intrinsic::aie2_upd_512_256_acc:
+  case Intrinsic::aie2_upd_1024_512_acc:
+  case Intrinsic::aie2_upd_1024_256_acc:
+    return VUpdateOpInfo{2, 3, 4};
+  default:
+    return std::nullopt;
+  }
+}
+
+std::optional<const AIEBaseInstrInfo::VExtractOpInfo>
+AIE2InstrInfo::getVExtractOpInfo(const MachineInstr &MI) const {
+
+  if (MI.getOpcode() != TargetOpcode::G_INTRINSIC)
+    return std::nullopt;
+
+  const GIntrinsic &GMI = cast<const GIntrinsic>(MI);
+
+  switch (GMI.getIntrinsicID()) {
+  case Intrinsic::aie2_ext_I256_I512:
+  case Intrinsic::aie2_ext_I512_I1024:
+  case Intrinsic::aie2_ext_I256_I1024:
+
+  case Intrinsic::aie2_ext_bf256_bf512:
+  case Intrinsic::aie2_ext_bf512_bf1024:
+  case Intrinsic::aie2_ext_bf256_bf1024:
+
+  case Intrinsic::aie2_ext_256_512_acc:
+  case Intrinsic::aie2_ext_512_1024_acc:
+  case Intrinsic::aie2_ext_256_1024_acc:
+    return VExtractOpInfo{2, 3};
+  default:
+    return std::nullopt;
+  }
+}

llvm/lib/Target/AIE/AIE2InstrInfo.h

@@ -169,6 +169,15 @@ class AIE2InstrInfo : public AIE2GenInstrInfo {
 
   bool canHoistCheapInst(const MachineInstr &MI) const override;
 
+  std::optional<const VConcatOpInfo>
+  getVConcatOpInfo(const MachineInstr &MI) const override;
+
+  std::optional<const VUpdateOpInfo>
+  getVUpdateOpInfo(const MachineInstr &MI) const override;
+
+  std::optional<const VExtractOpInfo>
+  getVExtractOpInfo(const MachineInstr &MI) const override;
+
 protected:
   SmallVector<AIEPseudoExpandInfo, 4>
   getSpillPseudoExpandInfo(const MachineInstr &MI) const override;

llvm/lib/Target/AIE/AIEBaseInstrInfo.cpp

@@ -924,3 +924,12 @@ const MIRFormatter *AIEBaseInstrInfo::getMIRFormatter() const {
     Formatter = std::make_unique<AIEMIRFormatter>();
   return Formatter.get();
 }
+
+/// Return operand information related to vector concat instrinsic.
+std::optional<const AIEBaseInstrInfo::VConcatOpInfo>
+AIEBaseInstrInfo::getVConcatOpInfo(const MachineInstr &MI) const {
+
+  if (MI.getOpcode() == TargetOpcode::G_CONCAT_VECTORS)
+    return VConcatOpInfo{1, 0};
+  return std::nullopt;
+}

llvm/lib/Target/AIE/AIEBaseInstrInfo.h

@@ -15,6 +15,7 @@
 #ifndef LLVM_LIB_TARGET_AIE_AIEBASEINSTRRINFO_H
 #define LLVM_LIB_TARGET_AIE_AIEBASEINSTRRINFO_H
 
+#include "AIE.h"
 #include "AIEMIRFormatter.h"
 #include "AIETiedRegOperands.h"
 #include "MCTargetDesc/AIEFormat.h"
@@ -379,6 +380,45 @@ struct AIEBaseInstrInfo : public TargetInstrInfo {
   static bool regClassMatches(const TargetRegisterClass &TRC,
                               const TargetRegisterClass *RC, unsigned Reg);
 
+  struct VConcatOpInfo {
+    // First input operand index.
+    unsigned FirstOperand;
+    // Number of non-register operands.
+    unsigned NumOfNonRegOperands;
+  };
+
+  /// Return operand information related to vector concat instrinsic.
+  virtual std::optional<const VConcatOpInfo>
+  getVConcatOpInfo(const MachineInstr &MI) const;
+
+  struct VUpdateOpInfo {
+    // Vector to update operand index.
+    unsigned Src;
+    // Subvector to insert.
+    unsigned SrcSubVec;
+    // Position to insert operand index.
+    unsigned SubVectorIndex;
+  };
+
+  /// Return operand information related to vector update instrinsic.
+  virtual std::optional<const VUpdateOpInfo>
+  getVUpdateOpInfo(const MachineInstr &MI) const {
+    llvm_unreachable("Target didn't implement getVUpdateOpInfo!");
+  }
+
+  struct VExtractOpInfo {
+    // Vector to update operand index.
+    unsigned Src;
+    // Position to extract.
+    unsigned SubVectorIndex;
+  };
+
+  /// Return operand information related to vector extract instrinsic.
+  virtual std::optional<const VExtractOpInfo>
+  getVExtractOpInfo(const MachineInstr &MI) const {
+    llvm_unreachable("Target didn't implement getVExtractOpInfo!");
+  }
+
 protected:
   /// Expand a spill pseudo-instruction into actual target instructions. This
   /// will essentially split the register being handled into its sub-registers,

llvm/lib/Target/AIE/AIECombine.td

@@ -70,6 +70,30 @@ def AIE2PostLegalizerGenericCombiner
     : GICombiner<"AIE2PostLegalizerGenericCombinerImpl", [ all_combines ]> {
 }
 
+def combine_extract_concat_matchdata: GIDefMatchData<"Register">;
+def combine_extract_concat : GICombineRule<
+  (defs root:$root, combine_extract_concat_matchdata:$matchinfo),
+  (match (wip_match_opcode G_INTRINSIC): $root,
+  [{ return matchExtractConcat(*${root}, MRI, (const AIEBaseInstrInfo &)B.getTII(), ${matchinfo}); }]),
+  (apply [{ applyExtractConcat(*${root}, MRI, B, ${matchinfo}); }])
+>;
+
+def combine_unmerge_concat_matchdata: GIDefMatchData<"std::pair<MachineInstr *, unsigned>">;
+def combine_unmerge_concat : GICombineRule<
+  (defs root:$root, combine_unmerge_concat_matchdata:$matchinfo),
+  (match (wip_match_opcode G_UNMERGE_VALUES): $root,
+  [{ return matchUnmergeConcat(*${root}, MRI, (const AIEBaseInstrInfo &)B.getTII(), ${matchinfo}); }]),
+  (apply [{ applyUnmergeConcat(*${root}, MRI, B, ${matchinfo}); }])
+>;
+
+def combine_upd_to_concat_matchdata: GIDefMatchData<"std::map<unsigned, Register>">;
+def combine_upd_to_concat : GICombineRule<
+  (defs root:$root, combine_upd_to_concat_matchdata:$matchinfo),
+  (match (wip_match_opcode G_INTRINSIC): $root,
+  [{ return matchUpdToConcat(*${root}, MRI, (const AIEBaseInstrInfo &)B.getTII(), ${matchinfo}); }]),
+  (apply [{ applyUpdToConcat(*${root}, MRI, B, ${matchinfo}); }])
+>;
+
 def load_store_increment_matchdata : GIDefMatchData<"AIELoadStoreCombineMatchData">;
 def combine_load_store_increment : GICombineRule <
   (defs root:$root, load_store_increment_matchdata:$matchinfo),
@@ -87,5 +111,9 @@ def combine_add_vector_elt_undef : GICombineRule <
 
 def AIE2PostLegalizerCustomCombiner
     : GICombiner<"AIE2PostLegalizerCustomCombinerImpl", [ combine_load_store_increment,
-                                                          combine_add_vector_elt_undef ]> {
+                                                          combine_add_vector_elt_undef,
+                                                          combine_extract_concat,
+                                                          combine_unmerge_concat,
+                                                          combine_upd_to_concat 
+                                                          ]> {
 }