Combine vextract and vbroadcast into VEXTBCST

llvm/lib/Target/AIE/AIE2PreLegalizerCombiner.cpp

@@ -73,6 +73,8 @@ class AIE2PreLegalizerCombinerImpl : public Combiner {
 
   bool tryToCombineVectorInserts(MachineInstr &MI, unsigned SclSrcBits) const;
 
+  bool tryToCombineExtBcst(MachineInstr &MI) const;
+
   bool tryToCombineIntrinsic(MachineInstr &MI) const;
 
 private:
@@ -241,6 +243,54 @@ bool AIE2PreLegalizerCombinerImpl::tryToCombineVectorInserts(
   return true;
 }
 
+// Combines vextract and vbroadcast into vextract_broadcast
+bool AIE2PreLegalizerCombinerImpl::tryToCombineExtBcst(MachineInstr &MI) const {
+  // Returns the combined intrinsicID for matching broadcast and extract ops
+  auto getExtBcstIntrinsicID = [](unsigned BcastID,
+                                  unsigned ExtID) -> std::optional<unsigned> {
+    switch (BcastID) {
+    case Intrinsic::aie2_vbroadcast8_I512:
+      if (ExtID == Intrinsic::aie2_vextract_elem8_I512)
+        return Intrinsic::aie2_vextract_broadcast8_I512;
+      break;
+    case Intrinsic::aie2_vbroadcast16_I512:
+      if (ExtID == Intrinsic::aie2_vextract_elem16_I512)
+        return Intrinsic::aie2_vextract_broadcast16_I512;
+      break;
+    case Intrinsic::aie2_vbroadcast32_I512:
+      if (ExtID == Intrinsic::aie2_vextract_elem32_I512)
+        return Intrinsic::aie2_vextract_broadcast32_I512;
+      break;
+    }
+    return std::nullopt;
+  };
+  assert(isa<GIntrinsic>(MI) && "this combine only supports instrinsics");
+  const Register DstReg = MI.getOperand(0).getReg();
+  MachineInstr *ExtMI = getDefIgnoringCopies(MI.getOperand(2).getReg(), MRI);
+  if (!isa<GIntrinsic>(*ExtMI))
+    return false;
+  // Checks for single use of extracted element
+  if (!MRI.hasOneNonDBGUse(ExtMI->getOperand(0).getReg()))
+    return false;
+
+  const unsigned BcstID = cast<GIntrinsic>(MI).getIntrinsicID();
+  const unsigned ExtID = cast<GIntrinsic>(*ExtMI).getIntrinsicID();
+  const std::optional<unsigned> ExtBcstIntrinsicID =
+      getExtBcstIntrinsicID(BcstID, ExtID);
+  if (!ExtBcstIntrinsicID)
+    return false;
+
+  const Register SrcReg = ExtMI->getOperand(2).getReg();
+  const Register IdxReg = ExtMI->getOperand(3).getReg();
+  MachineIRBuilder MIRBuilder(MI);
+  MIRBuilder.buildIntrinsic(*ExtBcstIntrinsicID, DstReg, false, false)
+      .addUse(SrcReg)
+      .addUse(IdxReg);
+  MI.eraseFromParent();
+
+  return true;
+}
+
 bool AIE2PreLegalizerCombinerImpl::tryToCombineIntrinsic(
     MachineInstr &MI) const {
   const unsigned IntrinsicID = cast<GIntrinsic>(MI).getIntrinsicID();
@@ -260,6 +310,12 @@ bool AIE2PreLegalizerCombinerImpl::tryToCombineIntrinsic(
   case Intrinsic::aie2_vinsert32_I512: {
     return tryToCombineVectorInserts(MI, getVInsertScalarSize(IntrinsicID));
   }
+  case Intrinsic::aie2_vbroadcast8_I512:
+  case Intrinsic::aie2_vbroadcast16_I512:
+  case Intrinsic::aie2_vbroadcast32_I512:
+  case Intrinsic::aie2_vbroadcast64_I512: {
+    return tryToCombineExtBcst(MI);
+  }
   default:
     break;
   }

llvm/test/CodeGen/AIE/aie2/GlobalISel/prelegalizercombiner-extract-broadcast.mir

@@ -0,0 +1,101 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+#
+# This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#
+# (c) Copyright 2025 Advanced Micro Devices, Inc. or its affiliates
+# RUN: llc -mtriple aie2 -run-pass=aie2-prelegalizer-combiner %s -verify-machineinstrs -o - | FileCheck %s
+
+---
+name:            vextract_broadcast8
+body:             |
+  bb.1.entry:
+    ; CHECK-LABEL: name: vextract_broadcast8
+    ; CHECK: [[COPY:%[0-9]+]]:_(<64 x s8>) = COPY $x2
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(s32) = COPY $r0
+    ; CHECK-NEXT: [[INT:%[0-9]+]]:_(<64 x s8>) = G_INTRINSIC intrinsic(@llvm.aie2.vextract.broadcast8.I512), [[COPY]](<64 x s8>), [[COPY1]](s32)
+    ; CHECK-NEXT: $x0 = COPY [[INT]](<64 x s8>)
+    %1:_(<64 x s8>) = COPY $x2
+    %2:_(s32) = COPY $r0
+    %4:_(s32) = G_CONSTANT i32 1
+    %3:_(s32) = G_INTRINSIC intrinsic(@llvm.aie2.vextract.elem8.I512), %1(<64 x s8>), %2(s32), %4(s32)
+    %0:_(<64 x s8>) = G_INTRINSIC intrinsic(@llvm.aie2.vbroadcast8.I512), %3(s32)
+    $x0 = COPY %0(<64 x s8>)
+...
+
+---
+name:            vextract_broadcast16
+body:             |
+  bb.1.entry:
+    ; CHECK-LABEL: name: vextract_broadcast16
+    ; CHECK: [[COPY:%[0-9]+]]:_(<32 x s16>) = COPY $x2
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(s32) = COPY $r0
+    ; CHECK-NEXT: [[INT:%[0-9]+]]:_(<32 x s16>) = G_INTRINSIC intrinsic(@llvm.aie2.vextract.broadcast16.I512), [[COPY]](<32 x s16>), [[COPY1]](s32)
+    ; CHECK-NEXT: $x0 = COPY [[INT]](<32 x s16>)
+    %1:_(<32 x s16>) = COPY $x2
+    %2:_(s32) = COPY $r0
+    %4:_(s32) = G_CONSTANT i32 1
+    %3:_(s32) = G_INTRINSIC intrinsic(@llvm.aie2.vextract.elem16.I512), %1(<32 x s16>), %2(s32), %4(s32)
+    %0:_(<32 x s16>) = G_INTRINSIC intrinsic(@llvm.aie2.vbroadcast16.I512), %3(s32)
+    $x0 = COPY %0(<32 x s16>)
+...
+
+---
+name:            vextract_broadcast32
+body:             |
+  bb.1.entry:
+    ; CHECK-LABEL: name: vextract_broadcast32
+    ; CHECK: [[COPY:%[0-9]+]]:_(<16 x s32>) = COPY $x2
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(s32) = COPY $r0
+    ; CHECK-NEXT: [[INT:%[0-9]+]]:_(<16 x s32>) = G_INTRINSIC intrinsic(@llvm.aie2.vextract.broadcast32.I512), [[COPY]](<16 x s32>), [[COPY1]](s32)
+    ; CHECK-NEXT: $x0 = COPY [[INT]](<16 x s32>)
+    %1:_(<16 x s32>) = COPY $x2
+    %2:_(s32) = COPY $r0
+    %4:_(s32) = G_CONSTANT i32 1
+    %3:_(s32) = G_INTRINSIC intrinsic(@llvm.aie2.vextract.elem32.I512), %1(<16 x s32>), %2(s32), %4(s32)
+    %0:_(<16 x s32>) = G_INTRINSIC intrinsic(@llvm.aie2.vbroadcast32.I512), %3(s32)
+    $x0 = COPY %0(<16 x s32>)
+...
+
+# Negative test case: Cannot be combined into vextract_broadcast due to type mismatch
+---
+name:            vextract_broadcast_type_mismatch
+body:             |
+  bb.1.entry:
+    ; CHECK-LABEL: name: vextract_broadcast_type_mismatch
+    ; CHECK: [[COPY:%[0-9]+]]:_(<32 x s16>) = COPY $x2
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(s32) = COPY $r0
+    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK-NEXT: [[INT:%[0-9]+]]:_(s32) = G_INTRINSIC intrinsic(@llvm.aie2.vextract.elem16.I512), [[COPY]](<32 x s16>), [[COPY1]](s32), [[C]](s32)
+    ; CHECK-NEXT: [[INT1:%[0-9]+]]:_(<16 x s32>) = G_INTRINSIC intrinsic(@llvm.aie2.vbroadcast32.I512), [[INT]](s32)
+    ; CHECK-NEXT: $x0 = COPY [[INT1]](<16 x s32>)
+    %1:_(<32 x s16>) = COPY $x2
+    %2:_(s32) = COPY $r0
+    %4:_(s32) = G_CONSTANT i32 1
+    %3:_(s32) = G_INTRINSIC intrinsic(@llvm.aie2.vextract.elem16.I512), %1(<32 x s16>), %2(s32), %4(s32)
+    %0:_(<16 x s32>) = G_INTRINSIC intrinsic(@llvm.aie2.vbroadcast32.I512), %3(s32)
+    $x0 = COPY %0(<16 x s32>)
+...
+
+# Negative test case: Cannot be combined into vextract_broadcast because extracted element has been used by instruction other than broadcast
+---
+name:            vextract_broadcast_multiuse
+body:             |
+  bb.1.entry:
+    ; CHECK-LABEL: name: vextract_broadcast_multiuse
+    ; CHECK: [[COPY:%[0-9]+]]:_(<32 x s16>) = COPY $x2
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(s32) = COPY $r0
+    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK-NEXT: [[INT:%[0-9]+]]:_(s32) = G_INTRINSIC intrinsic(@llvm.aie2.vextract.elem16.I512), [[COPY]](<32 x s16>), [[COPY1]](s32), [[C]](s32)
+    ; CHECK-NEXT: [[INT1:%[0-9]+]]:_(<32 x s16>) = G_INTRINSIC intrinsic(@llvm.aie2.vbroadcast16.I512), [[INT]](s32)
+    ; CHECK-NEXT: $x0 = COPY [[INT1]](<32 x s16>)
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[INT]](s32)
+    %1:_(<32 x s16>) = COPY $x2
+    %2:_(s32) = COPY $r0
+    %4:_(s32) = G_CONSTANT i32 1
+    %3:_(s32) = G_INTRINSIC intrinsic(@llvm.aie2.vextract.elem16.I512), %1(<32 x s16>), %2(s32), %4(s32)
+    %0:_(<32 x s16>) = G_INTRINSIC intrinsic(@llvm.aie2.vbroadcast16.I512), %3(s32)
+    $x0 = COPY %0(<32 x s16>)
+    PseudoRET implicit $lr, implicit %3
+...