[AIE] Combine REG_SEQ of copies of subregs into a copy of the superreg

Martien de Jong · martien-de-jong · commit 598be5c80745 · 2025-04-30T17:54:09.000+02:00
diff --git a/llvm/lib/Target/AIE/AIEPostSelectOptimize.cpp b/llvm/lib/Target/AIE/AIEPostSelectOptimize.cpp
@@ -617,6 +617,73 @@ bool modifyStoreFlush(MachineBasicBlock &MBB, MachineRegisterInfo &MRI) {
   return Changed;
 }
 
+static bool combineSubRegCopyToSuperregCopy(MachineBasicBlock &MBB,
+                                            MachineRegisterInfo &MRI) {
+  const TargetInstrInfo *TII = MBB.getParent()->getSubtarget().getInstrInfo();
+  bool Changed = false;
+
+  for (MachineInstr &MI : make_early_inc_range(reverse(MBB))) {
+    if (!MI.isRegSequence())
+      continue;
+
+    LLVM_DEBUG(llvm::dbgs() << "SubregCopyToSuperRegCopy visiting ");
+    LLVM_DEBUG(MI.dump());
+    Register CommonSrcReg;
+    // TODO: We may want to check that the list covers the whole superregister
+    // lest we run the risk of replacing an undefined register with the defined
+    // source subreg.
+    for (unsigned I = 1; I < MI.getNumExplicitOperands(); I += 2) {
+
+      const Register SrcReg = MI.getOperand(I).getReg();
+      const unsigned SubRegIdx = MI.getOperand(I + 1).getImm();
+
+      auto *SubRegDef = MRI.getVRegDef(SrcReg);
+      if (!SubRegDef->isCopy()) {
+        LLVM_DEBUG(llvm::dbgs() << "Src is not a COPY. Skip\n");
+        CommonSrcReg = {};
+        break;
+      }
+
+      if (!SubRegDef->getOperand(1).getSubReg()) {
+        LLVM_DEBUG(llvm::dbgs() << "Src COPY is not a subreg COPY. Skip\n");
+        CommonSrcReg = {};
+        break;
+      }
+
+      const unsigned SrcSubRegIdx = SubRegDef->getOperand(1).getSubReg();
+
+      if (SubRegIdx != SrcSubRegIdx) {
+        LLVM_DEBUG(llvm::dbgs() << "SubReg indexes are not the same. Skip\n");
+        CommonSrcReg = {};
+        break;
+      }
+
+      Register SubRegCopySrc = SubRegDef->getOperand(1).getReg();
+      if (CommonSrcReg && SubRegCopySrc != CommonSrcReg) {
+        LLVM_DEBUG(llvm::dbgs()
+                   << "REG_SEQUENCE from different sources. Skip\n");
+        CommonSrcReg = {};
+        break;
+      }
+      CommonSrcReg = SubRegCopySrc;
+    }
+
+    if (CommonSrcReg) {
+      LLVM_DEBUG(llvm::dbgs() << "Folding ");
+      LLVM_DEBUG(MI.dump());
+      LLVM_DEBUG(llvm::dbgs() << "into superreg COPY\n");
+      auto MIB =
+          BuildMI(MBB, MI, MI.getDebugLoc(), TII->get(TargetOpcode::COPY),
+                  MI.getOperand(0).getReg())
+              .addReg(CommonSrcReg);
+      LLVM_DEBUG(MIB->dump());
+      MI.eraseFromParent();
+      Changed = true;
+    }
+  }
+  return Changed;
+}
+
 bool AIEPostSelectOptimize::runOnMachineFunction(MachineFunction &MF) {
   LLVM_DEBUG(dbgs() << "\n******* POST I-SEL OPTIMIZATION PASS *******\n"
                     << "********** Function: " << MF.getName() << '\n');
@@ -632,6 +699,12 @@ bool AIEPostSelectOptimize::runOnMachineFunction(MachineFunction &MF) {
 
   bool Changed = false;
 
+  // 0. Fold REG_SEQUENCE (COPY %0.sub_bfp16_x), %subreg.sub_bfp16_x,
+  // (%0.sub_bfp16_e), %subreg.sub_bfp16_e) into COPY %0
+  for (MachineBasicBlock &MBB : MF) {
+    Changed |= combineSubRegCopyToSuperregCopy(MBB, MF.getRegInfo());
+  }
+
   // 1. Turn INSERT_SUBREG into REG_SEQUENCE when possible
   for (MachineBasicBlock &MBB : MF) {
     Changed |= combineINSERT_SUBREG(MBB);
diff --git a/llvm/test/CodeGen/AIE/aie2/end-to-end/TanhTemplated-swp.ll b/llvm/test/CodeGen/AIE/aie2/end-to-end/TanhTemplated-swp.ll
@@ -79,18 +79,18 @@ define dso_local void @TanhTemplated(ptr noalias %ifm, ptr noalias %ofm, ptr non
 ; CHECK-NEXT:    movxm r6, #15616; vmul.f bmh2, x0, x3, r1
 ; CHECK-NEXT:    movxm r7, #16000
 ; CHECK-NEXT:    vbcst.16 x1, r3
-; CHECK-NEXT:    vbcst.16 x10, r4
-; CHECK-NEXT:    vbcst.16 x8, r5; vmul.f bmh3, x0, x3, r1
+; CHECK-NEXT:    vbcst.16 x8, r4
+; CHECK-NEXT:    vbcst.16 x10, r5; vmul.f bmh3, x0, x3, r1
 ; CHECK-NEXT:    vbcst.16 x6, r6
 ; CHECK-NEXT:    vconv.bf16.fp32 wl3, bmh2; vbcst.16 x4, r7
 ; CHECK-NEXT:    vmov wh6, wl2
 ; CHECK-NEXT:    vmin_ge.bf16 x3, r16, x3, x1
-; CHECK-NEXT:    or r8, r16, r16; vmax_lt.bf16 x3, r16, x3, x10
-; CHECK-NEXT:    vconv.bf16.fp32 wl5, bmh3; vband x7, x8, x3
+; CHECK-NEXT:    or r8, r16, r16; vmax_lt.bf16 x3, r16, x3, x8
+; CHECK-NEXT:    vconv.bf16.fp32 wl5, bmh3; vband x7, x10, x3
 ; CHECK-NEXT:    vmov wh7, wl2
 ; CHECK-NEXT:    vmin_ge.bf16 x5, r16, x5, x1
-; CHECK-NEXT:    vldb wl7, [p0], #32; vmax_lt.bf16 x5, r16, x5, x10
-; CHECK-NEXT:    vband x7, x8, x5
+; CHECK-NEXT:    vldb wl7, [p0], #32; vmax_lt.bf16 x5, r16, x5, x8
+; CHECK-NEXT:    vband x7, x10, x5
 ; CHECK-NEXT:    vldb wl7, [p0], #32; vmov wh7, wl2; vmul.f bmh2, x6, x7, r1
 ; CHECK-NEXT:    vmov wh4, wl2
 ; CHECK-NEXT:    vmov wh3, wl2; vmul.f bmh4, x6, x7, r1
@@ -105,16 +105,16 @@ define dso_local void @TanhTemplated(ptr noalias %ifm, ptr noalias %ofm, ptr non
 ; CHECK-NEXT:    vconv.bf16.fp32 wl3, bmh5; movxm le, #.L_LEnd0
 ; CHECK-NEXT:    add.nc lc, r2, #-2
 ; CHECK-NEXT:    vconv.bf16.fp32 wl5, bmh7; vmin_ge.bf16 x3, r16, x3, x1
-; CHECK-NEXT:    vmax_lt.bf16 x3, r16, x3, x10
+; CHECK-NEXT:    vmax_lt.bf16 x3, r16, x3, x8
 ; CHECK-NEXT:    mova r0, #28; vconv.bf16.fp32 wl7, bmh3; vmin_ge.bf16 x11, r16, x5, x1
 ; CHECK-NEXT:    .p2align 4
 ; CHECK-NEXT:  .LBB0_1: // %for.body
 ; CHECK-NEXT:    // =>This Inner Loop Header: Depth=1
-; CHECK-NEXT:    nopb ; nopa ; nops ; nopx ; vband x9, x8, x3; nopv
+; CHECK-NEXT:    nopb ; nopa ; nops ; nopx ; vband x9, x10, x3; nopv
 ; CHECK-NEXT:    vldb wl7, [p0], #32; vmov wh3, wl2
 ; CHECK-NEXT:    nopx ; vmov wh9, wl2; vmul.f bmh5, x7, x0, r1
-; CHECK-NEXT:    vconv.bf16.fp32 wl7, bml4; vldb wl7, [p0], #32; vmax_lt.bf16 x5, r16, x11, x10; vmac.f bmh4, bmh0, x3, x4, r1
-; CHECK-NEXT:    vband x9, x8, x5; vmul.f bmh2, x6, x9, r1
+; CHECK-NEXT:    vconv.bf16.fp32 wl7, bml4; vldb wl7, [p0], #32; vmax_lt.bf16 x5, r16, x11, x8; vmac.f bmh4, bmh0, x3, x4, r1
+; CHECK-NEXT:    vband x9, x10, x5; vmul.f bmh2, x6, x9, r1
 ; CHECK-NEXT:    vmov wh9, wl2; vmul.f bmh6, x7, x0, r1
 ; CHECK-NEXT:    vsub.f bml0, bmh5, bmh1, r0
 ; CHECK-NEXT:    vmul.f bmh3, x6, x9, r1
@@ -127,29 +127,29 @@ define dso_local void @TanhTemplated(ptr noalias %ifm, ptr noalias %ofm, ptr non
 ; CHECK-NEXT:    vconv.bf16.fp32 wl3, bmh7
 ; CHECK-NEXT:    vst.conv.bf16.fp32 bml0, [p1], #32; vmsc.f bml4, bml2, x3, x5, r1
 ; CHECK-NEXT:    vconv.bf16.fp32 wl5, bmh8; vmin_ge.bf16 x9, r16, x3, x1
-; CHECK-NEXT:    vst.conv.bf16.fp32 bml1, [p1], #32; vmax_lt.bf16 x3, r16, x9, x10
+; CHECK-NEXT:    vst.conv.bf16.fp32 bml1, [p1], #32; vmax_lt.bf16 x3, r16, x9, x8
 ; CHECK-NEXT:  .L_LEnd0:
 ; CHECK-NEXT:    nopb ; nopa ; vconv.bf16.fp32 wl7, bml3; nopx ; vmin_ge.bf16 x11, r16, x5, x1; nopv
 ; CHECK-NEXT:  // %bb.2:
-; CHECK-NEXT:    nopa ; nopb ; nopxm
+; CHECK-NEXT:    nopb ; nopa ; nops ; nopxm ; nopv
 ; CHECK-NEXT:    vmov wh7, wl2
 ; CHECK-NEXT:    vconv.bf16.fp32 wl1, bml4; vmov wh1, wl2
-; CHECK-NEXT:    vmov wh6, wl2; vmul.f bmh3, x7, x0, r1
-; CHECK-NEXT:    vmax_lt.bf16 x10, r16, x11, x10; vmul.f bmh2, x1, x0, r1
-; CHECK-NEXT:    vband x1, x8, x3
-; CHECK-NEXT:    vband x8, x8, x10
+; CHECK-NEXT:    vmul.f bmh3, x7, x0, r1
+; CHECK-NEXT:    vmax_lt.bf16 x8, r16, x11, x8; vmul.f bmh2, x1, x0, r1
+; CHECK-NEXT:    vband x1, x10, x3
+; CHECK-NEXT:    vband x10, x10, x8
 ; CHECK-NEXT:    vmov wh1, wl2; vsub.f bmh3, bmh3, bmh1, r0
-; CHECK-NEXT:    vmov wh8, wl2; vsub.f bmh2, bmh2, bmh1, r0
+; CHECK-NEXT:    vmov wh10, wl2; vsub.f bmh2, bmh2, bmh1, r0
 ; CHECK-NEXT:    vmul.f bmh2, x6, x1, r1
-; CHECK-NEXT:    vmov wh4, wl2; vmul.f bmh3, x6, x8, r1
+; CHECK-NEXT:    vmul.f bmh3, x6, x10, r1
 ; CHECK-NEXT:    vmov wh3, wl2
-; CHECK-NEXT:    vmov wh10, wl2
+; CHECK-NEXT:    vmov wh8, wl2
 ; CHECK-NEXT:    vst.conv.bf16.fp32 bmh3, [p1], #32; vmac.f bmh4, bmh0, x3, x4, r1
-; CHECK-NEXT:    vst.conv.bf16.fp32 bmh2, [p1], #32; vmac.f bmh0, bmh0, x10, x4, r1
-; CHECK-NEXT:    vconv.bf16.fp32 wl4, bmh2
+; CHECK-NEXT:    vst.conv.bf16.fp32 bmh2, [p1], #32; vmac.f bmh0, bmh0, x8, x4, r1
+; CHECK-NEXT:    vconv.bf16.fp32 wl4, bmh2; vmov wh4, wl2
 ; CHECK-NEXT:    vconv.bf16.fp32 wl4, bmh3
 ; CHECK-NEXT:    vmsc.f bmh2, bmh4, x4, x3, r1
-; CHECK-NEXT:    vmsc.f bmh0, bmh0, x4, x10, r1
+; CHECK-NEXT:    vmsc.f bmh0, bmh0, x4, x8, r1
 ; CHECK-NEXT:    nop
 ; CHECK-NEXT:    nop
 ; CHECK-NEXT:    nop
diff --git a/llvm/test/CodeGen/AIE/aie2p/GlobalIsel/post-select-combine.mir b/llvm/test/CodeGen/AIE/aie2p/GlobalIsel/post-select-combine.mir
@@ -0,0 +1,167 @@
+# 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 aie2p -run-pass=aie-post-select-optimize,dead-mi-elimination %s -verify-machineinstrs -o - | FileCheck %s
+
+# This file tests various post-select combines for aie2p
+
+---
+name:            REG_SEQ_of_subreg_copies
+alignment:       16
+legalized:       true
+regBankSelected: true
+selected:        true
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+  liveins: $ex0
+    ; CHECK-LABEL: name: REG_SEQ_of_subreg_copies
+    ; CHECK: liveins: $ex0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:vec576 = COPY $ex0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:vec576 = COPY [[COPY]]
+    ; CHECK-NEXT: $ex0 = COPY [[COPY1]]
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit $ex0
+    %0:vec576 = COPY $ex0
+    %1:vec512 = COPY %0.sub_bfp16_x
+    %2:expvec64 = COPY %0.sub_bfp16_e
+    %3:vec576 = REG_SEQUENCE %1, %subreg.sub_bfp16_x, %2, %subreg.sub_bfp16_e
+    $ex0 = COPY %3
+    PseudoRET implicit $lr, implicit $ex0
+...
+
+---
+name:            dont_REG_SEQ_of_different_sources
+alignment:       16
+legalized:       true
+regBankSelected: true
+selected:        true
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+  liveins: $ex0, $ex1
+    ; CHECK-LABEL: name: dont_REG_SEQ_of_different_sources
+    ; CHECK: liveins: $ex0, $ex1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:vec576 = COPY $ex0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:vec576 = COPY $ex1
+    ; CHECK-NEXT: [[COPY2:%[0-9]+]]:vec512 = COPY [[COPY]].sub_bfp16_x
+    ; CHECK-NEXT: [[COPY3:%[0-9]+]]:expvec64 = COPY [[COPY1]].sub_bfp16_e
+    ; CHECK-NEXT: [[REG_SEQUENCE:%[0-9]+]]:vec576 = REG_SEQUENCE [[COPY2]], %subreg.sub_bfp16_x, [[COPY3]], %subreg.sub_bfp16_e
+    ; CHECK-NEXT: $ex0 = COPY [[REG_SEQUENCE]]
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit $ex0
+    %0:vec576 = COPY $ex0
+    %1:vec576 = COPY $ex1
+    %2:vec512 = COPY %0.sub_bfp16_x
+    %3:expvec64 = COPY %1.sub_bfp16_e
+    %4:vec576 = REG_SEQUENCE %2, %subreg.sub_bfp16_x, %3, %subreg.sub_bfp16_e
+    $ex0 = COPY %4
+    PseudoRET implicit $lr, implicit $ex0
+...
+---
+name:            dont_REG_SEQ_of_non_subreg
+alignment:       16
+legalized:       true
+regBankSelected: true
+selected:        true
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+  liveins: $ex0
+    ; CHECK-LABEL: name: dont_REG_SEQ_of_non_subreg
+    ; CHECK: liveins: $ex0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:vec576 = COPY $ex0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:vec512 = COPY [[COPY]].sub_bfp16_x
+    ; CHECK-NEXT: [[COPY2:%[0-9]+]]:expvec64 = COPY $e0
+    ; CHECK-NEXT: [[REG_SEQUENCE:%[0-9]+]]:vec576 = REG_SEQUENCE [[COPY1]], %subreg.sub_bfp16_x, [[COPY2]], %subreg.sub_bfp16_e
+    ; CHECK-NEXT: $ex0 = COPY [[REG_SEQUENCE]]
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit $ex0
+    %0:vec576 = COPY $ex0
+    %2:vec512 = COPY %0.sub_bfp16_x
+    %3:expvec64 = COPY $e0
+    %4:vec576 = REG_SEQUENCE %2, %subreg.sub_bfp16_x, %3, %subreg.sub_bfp16_e
+    $ex0 = COPY %4
+    PseudoRET implicit $lr, implicit $ex0
+...
+---
+name:            do_REG_SEQ_of_subregs
+alignment:       16
+legalized:       true
+regBankSelected: true
+selected:        true
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+  liveins: $x0
+    ; CHECK-LABEL: name: do_REG_SEQ_of_subregs
+    ; CHECK: liveins: $x0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:vec512 = COPY $x0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:vec512 = COPY [[COPY]]
+    ; CHECK-NEXT: $x0 = COPY [[COPY1]]
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit $x0
+    %0:vec512 = COPY $x0
+    %2:vec256 = COPY %0.sub_256_lo
+    %3:vec256 = COPY %0.sub_256_hi
+    %4:vec512 = REG_SEQUENCE %2, %subreg.sub_256_lo, %3, %subreg.sub_256_hi
+    $x0 = COPY %4
+    PseudoRET implicit $lr, implicit $x0
+...
+---
+name:            dont_REG_SEQ_of_swapped_index
+alignment:       16
+legalized:       true
+regBankSelected: true
+selected:        true
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+  liveins: $x0
+    ; CHECK-LABEL: name: dont_REG_SEQ_of_swapped_index
+    ; CHECK: liveins: $x0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:vec512 = COPY $x0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:vec256 = COPY [[COPY]].sub_256_lo
+    ; CHECK-NEXT: [[COPY2:%[0-9]+]]:vec256 = COPY [[COPY]].sub_256_hi
+    ; CHECK-NEXT: [[REG_SEQUENCE:%[0-9]+]]:vec512 = REG_SEQUENCE [[COPY2]], %subreg.sub_256_lo, [[COPY1]], %subreg.sub_256_hi
+    ; CHECK-NEXT: $x0 = COPY [[REG_SEQUENCE]]
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit $x0
+    %0:vec512 = COPY $x0
+    %2:vec256 = COPY %0.sub_256_lo
+    %3:vec256 = COPY %0.sub_256_hi
+    %4:vec512 = REG_SEQUENCE %3, %subreg.sub_256_lo, %2, %subreg.sub_256_hi
+    $x0 = COPY %4
+    PseudoRET implicit $lr, implicit $x0
+...
+---
+name:            dont_REG_SEQ_of_equal_index
+alignment:       16
+legalized:       true
+regBankSelected: true
+selected:        true
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+  liveins: $x0
+    ; CHECK-LABEL: name: dont_REG_SEQ_of_equal_index
+    ; CHECK: liveins: $x0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:vec512 = COPY $x0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:vec256 = COPY [[COPY]].sub_256_lo
+    ; CHECK-NEXT: [[COPY2:%[0-9]+]]:vec256 = COPY [[COPY]].sub_256_lo
+    ; CHECK-NEXT: [[REG_SEQUENCE:%[0-9]+]]:vec512 = REG_SEQUENCE [[COPY1]], %subreg.sub_256_lo, [[COPY2]], %subreg.sub_256_hi
+    ; CHECK-NEXT: $x0 = COPY [[REG_SEQUENCE]]
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit $x0
+    %0:vec512 = COPY $x0
+    %2:vec256 = COPY %0.sub_256_lo
+    %3:vec256 = COPY %0.sub_256_lo
+    %4:vec512 = REG_SEQUENCE %2, %subreg.sub_256_lo, %3, %subreg.sub_256_hi
+    $x0 = COPY %4
+    PseudoRET implicit $lr, implicit $x0
+...
