From 972bc8b0322a8b783d33b88e0e553a70caa2c2cb Mon Sep 17 00:00:00 2001
From: Sai Abhinay Anubola <saiabhinay.anubola@amd.com>
Date: Thu, 21 Nov 2024 12:17:03 +0530
Subject: [PATCH] Add support for bfloat in SetExtract PreLegalizerCombiner

---
 .../Target/AIE/AIE2PreLegalizerCombiner.cpp   | 39 ++++++++----
 .../combine-set-extract-prelegalizer.mir      | 60 +++++++++++++++++++
 2 files changed, 86 insertions(+), 13 deletions(-)

diff --git a/llvm/lib/Target/AIE/AIE2PreLegalizerCombiner.cpp b/llvm/lib/Target/AIE/AIE2PreLegalizerCombiner.cpp
index 500b06adf019..edd7723ba328 100644
--- a/llvm/lib/Target/AIE/AIE2PreLegalizerCombiner.cpp
+++ b/llvm/lib/Target/AIE/AIE2PreLegalizerCombiner.cpp
@@ -149,6 +149,18 @@ bool AIE2PreLegalizerCombinerImpl::tryToCombineVectorShiftsByZero(
 
 bool AIE2PreLegalizerCombinerImpl::tryToCombineSetExtract(
     MachineInstr &MI) const {
+  // Compares the index registers and their associated constant values from two
+  // given machine instructions.
+  auto IsEqualIdx = [&](const MachineInstr *SetMI, const MachineInstr *ExtMI) {
+    const Register SetMIIdxReg = SetMI->getOperand(3).getReg();
+    const Register ExtMIIdxReg = ExtMI->getOperand(3).getReg();
+    if (SetMIIdxReg == ExtMIIdxReg)
+      return true;
+    auto SetMICst = getIConstantVRegValWithLookThrough(SetMIIdxReg, MRI);
+    auto ExtMICst = getIConstantVRegValWithLookThrough(ExtMIIdxReg, MRI);
+    return (SetMICst && ExtMICst &&
+            SetMICst->Value.getZExtValue() == ExtMICst->Value.getZExtValue());
+  };
   const Register DstReg = MI.getOperand(0).getReg();
   MachineInstr *ExtOp = getDefIgnoringCopies(MI.getOperand(2).getReg(), MRI);
 
@@ -162,18 +174,18 @@ bool AIE2PreLegalizerCombinerImpl::tryToCombineSetExtract(
     break;
   }
   case Intrinsic::aie2_set_I512_I256: {
-    if (cast<GIntrinsic>(*ExtOp).getIntrinsicID() !=
-        Intrinsic::aie2_ext_I256_I512)
-      return false;
-    const Register SetOpIdxReg = MI.getOperand(3).getReg();
-    const Register ExtOpIdxReg = ExtOp->getOperand(3).getReg();
-    auto SetOpCst = getIConstantVRegValWithLookThrough(SetOpIdxReg, MRI);
-    auto ExtOpCst = getIConstantVRegValWithLookThrough(ExtOpIdxReg, MRI);
-    if (SetOpIdxReg != ExtOpIdxReg &&
-        (!SetOpCst || !ExtOpCst ||
-         SetOpCst->Value.getZExtValue() != ExtOpCst->Value.getZExtValue()))
-      return false;
-    break;
+    if (cast<GIntrinsic>(*ExtOp).getIntrinsicID() ==
+            Intrinsic::aie2_ext_I256_I512 &&
+        IsEqualIdx(&MI, ExtOp))
+      break;
+    return false;
+  }
+  case Intrinsic::aie2_set_bf512_bf256: {
+    if (cast<GIntrinsic>(*ExtOp).getIntrinsicID() ==
+            Intrinsic::aie2_ext_bf256_bf512 &&
+        IsEqualIdx(&MI, ExtOp))
+      break;
+    return false;
   }
   default:
     return false;
@@ -298,7 +310,8 @@ bool AIE2PreLegalizerCombinerImpl::tryToCombineIntrinsic(
   case Intrinsic::aie2_set_I512_I128: {
     return tryToCombineSetExtract(MI);
   }
-  case Intrinsic::aie2_set_I512_I256: {
+  case Intrinsic::aie2_set_I512_I256:
+  case Intrinsic::aie2_set_bf512_bf256: {
     return Combine256To512SetExtract && tryToCombineSetExtract(MI);
   }
   case Intrinsic::aie2_vinsert8_I512:
diff --git a/llvm/test/CodeGen/AIE/aie2/GlobalISel/combine-set-extract-prelegalizer.mir b/llvm/test/CodeGen/AIE/aie2/GlobalISel/combine-set-extract-prelegalizer.mir
index 05f5239a1671..675bbc72c2e9 100644
--- a/llvm/test/CodeGen/AIE/aie2/GlobalISel/combine-set-extract-prelegalizer.mir
+++ b/llvm/test/CodeGen/AIE/aie2/GlobalISel/combine-set-extract-prelegalizer.mir
@@ -83,3 +83,63 @@ body:             |
     %3:_(<16 x s32>) = G_INTRINSIC intrinsic(@llvm.aie2.set.I512.I256), %2(<8 x s32>), %10(s32)
     PseudoRET implicit $lr, implicit %3
 ...
+
+---
+name:            set-extract-bf256
+legalized:       false
+body:             |
+  bb.1.entry:
+    liveins: $x0
+    ; CHECK-LABEL: name: set-extract-bf256
+    ; CHECK: liveins: $x0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<16 x s32>) = COPY $x0
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[COPY]](<16 x s32>)
+    %0:_(<16 x s32>) = COPY $x0
+    %1:_(s32) = G_CONSTANT i32 0
+    %2:_(<8 x s32>) = G_INTRINSIC intrinsic(@llvm.aie2.ext.bf256.bf512), %0(<16 x s32>), %1(s32)
+    %3:_(<16 x s32>) = G_INTRINSIC intrinsic(@llvm.aie2.set.bf512.bf256), %2(<8 x s32>), %1(s32)
+    PseudoRET implicit $lr, implicit %3
+...
+
+---
+name:            set-extract-bf256_different_reg_same_idx
+legalized:       false
+body:             |
+  bb.1.entry:
+    liveins: $x0
+    ; CHECK-LABEL: name: set-extract-bf256_different_reg_same_idx
+    ; CHECK: liveins: $x0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<16 x s32>) = COPY $x0
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[COPY]](<16 x s32>)
+    %0:_(<16 x s32>) = COPY $x0
+    %1:_(s32) = G_CONSTANT i32 0
+    %10:_(s32) = G_CONSTANT i32 0
+    %2:_(<8 x s32>) = G_INTRINSIC intrinsic(@llvm.aie2.ext.bf256.bf512), %0(<16 x s32>), %1(s32)
+    %3:_(<16 x s32>) = G_INTRINSIC intrinsic(@llvm.aie2.set.bf512.bf256), %2(<8 x s32>), %10(s32)
+    PseudoRET implicit $lr, implicit %3
+...
+
+---
+name:            set-extract-bf256_different_idx
+legalized:       false
+body:             |
+  bb.1.entry:
+    liveins: $x0
+    ; CHECK-LABEL: name: set-extract-bf256_different_idx
+    ; CHECK: liveins: $x0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<16 x s32>) = COPY $x0
+    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
+    ; CHECK-NEXT: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK-NEXT: [[INT:%[0-9]+]]:_(<8 x s32>) = G_INTRINSIC intrinsic(@llvm.aie2.ext.bf256.bf512), [[COPY]](<16 x s32>), [[C]](s32)
+    ; CHECK-NEXT: [[INT1:%[0-9]+]]:_(<16 x s32>) = G_INTRINSIC intrinsic(@llvm.aie2.set.bf512.bf256), [[INT]](<8 x s32>), [[C1]](s32)
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[INT1]](<16 x s32>)
+    %0:_(<16 x s32>) = COPY $x0
+    %1:_(s32) = G_CONSTANT i32 0
+    %10:_(s32) = G_CONSTANT i32 1
+    %2:_(<8 x s32>) = G_INTRINSIC intrinsic(@llvm.aie2.ext.bf256.bf512), %0(<16 x s32>), %1(s32)
+    %3:_(<16 x s32>) = G_INTRINSIC intrinsic(@llvm.aie2.set.bf512.bf256), %2(<8 x s32>), %10(s32)
+    PseudoRET implicit $lr, implicit %3
+...