Skip to content

Commit

Permalink
[AIE2P] Legalize G_UNMERGE_VALUES 256-bit vector to 2 128-bit vectors
Browse files Browse the repository at this point in the history
  • Loading branch information
@niwinanto
niwinanto committed Jan 21, 2025
1 parent 739dea5 commit d6e46e6
Show file tree
Hide file tree
Showing 4 changed files with 122 additions and 3 deletions.
62 changes: 62 additions & 0 deletions llvm/lib/Target/AIE/AIELegalizerHelper.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -231,6 +231,63 @@ bool AIELegalizerHelper::legalizeG_BUILD_VECTOR(LegalizerHelper &Helper,
return true;
}

// Legalize G_UNMERGE_VALUES of 256-bit vector into 2 x 128-bit vector. For an
// example,
// %1:_(<4 x s32>), %2:_(<4 x s32>) = G_UNMERGE_VALUES %0:_(<8 x s32>)
// to
// %1:_(<4 x s32>) = G_AIE_UNPAD_VECTOR %0(<8 x s32>)
// %3:_(s32) = G_CONSTANT i32 16
// %4:_(<16 x s32>) = G_CONCAT_VECTORS %0(<8 x s32>), %0(<8 x s32>)
// %5:_(<16 x s32>) = G_IMPLICIT_DEF
// %6:_(<16 x s32>) = G_AIE_VSHIFT_RIGHT %4, %5, %3(s32)
// %2:_(<4 x s32>) = G_AIE_UNPAD_VECTOR %6(<16 x s32>)
bool AIELegalizerHelper::legalizeG_UNMERGE_VALUES_128bit(
LegalizerHelper &Helper, MachineInstr &MI) const {
MachineIRBuilder &MIRBuilder = Helper.MIRBuilder;
MachineRegisterInfo &MRI = *MIRBuilder.getMRI();

assert(MI.getNumOperands() == 3 &&
"Expected G_UNMERGE_VALUES of 256-bit vector to 2 x 128-bit vector");
const AIEBaseInstrInfo *II = ST.getInstrInfo();
const unsigned UnpadOpc = II->getGenericUnpadVectorOpcode();
const Register SrcReg = MI.getOperand(MI.getNumOperands() - 1).getReg();
const Register DstRegLow = MI.getOperand(0).getReg();
const Register DstRegHigh = MI.getOperand(1).getReg();
const LLT SrcTy = MRI.getType(SrcReg);

// Extracting lower 128-bit is easy: just discard (unpad) the high bits
MIRBuilder.buildInstr(UnpadOpc, {DstRegLow}, {SrcReg});

// To extract the higher 128-bit, we need to shift them to the lower position,
// then unpad again.
// We need to shift the upper 128-bit content by 16-byte (128-bit)
auto ShiftAmt = MIRBuilder.buildConstant(LLT::scalar(32), 16);

// VSHIFT operates on 512-bit inputs. We need to pad the 256-bit source
// operand to 512-bit
const Register ImplicitDef256 = MIRBuilder.buildUndef(SrcTy).getReg(0);

const LLT Vec512 = SrcTy.multiplyElements(2);

// Create the first 512-bit vector input
auto ConcatValue =
MIRBuilder.buildConcatVectors({Vec512}, {SrcReg, ImplicitDef256});

// The second input will be ignored. Just create a dummy input
auto ImplicitDef512 = MIRBuilder.buildUndef(Vec512);

// Now create the VSHIFT
const unsigned VShiftOpc = II->getGenericVShiftOpcode();
auto VShift = MIRBuilder.buildInstr(VShiftOpc, {Vec512},
{ConcatValue, ImplicitDef512, ShiftAmt});

// Finally, unpad the 512-bit result to 128-bit
MIRBuilder.buildInstr(UnpadOpc, {DstRegHigh}, {VShift});

MI.eraseFromParent();
return true;
}

bool AIELegalizerHelper::legalizeG_UNMERGE_VALUES(LegalizerHelper &Helper,
MachineInstr &MI) const {
MachineIRBuilder &MIRBuilder = Helper.MIRBuilder;
Expand All @@ -241,6 +298,11 @@ bool AIELegalizerHelper::legalizeG_UNMERGE_VALUES(LegalizerHelper &Helper,
const Register LastReg = MI.getOperand(MI.getNumOperands() - 1).getReg();
const LLT FirstTy = MRI.getType(FirstReg);
const LLT LastTy = MRI.getType(LastReg);

if (ST.isAIE2P() && FirstTy.isVector() && FirstTy.getSizeInBits() == 128 &&
LastTy.getSizeInBits() == 256)
return legalizeG_UNMERGE_VALUES_128bit(Helper, MI);

assert(LastTy.isVector() &&
(FirstTy.getScalarSizeInBits() * (MI.getNumOperands() - 1)) ==
LastTy.getSizeInBits() &&
Expand Down
2 changes: 2 additions & 0 deletions llvm/lib/Target/AIE/AIELegalizerHelper.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,8 @@ class AIELegalizerHelper {
Register SourceReg) const;
bool unpack32BitVector(LegalizerHelper &Helper, MachineInstr &MI,
Register SourceReg) const;
bool legalizeG_UNMERGE_VALUES_128bit(LegalizerHelper &Helper,
MachineInstr &MI) const;
};

} // namespace llvm
Expand Down
10 changes: 7 additions & 3 deletions llvm/lib/Target/AIE/aie2p/AIE2PLegalizerInfo.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -564,9 +564,13 @@ AIE2PLegalizerInfo::AIE2PLegalizerInfo(const AIE2PSubtarget &ST)
.customIf([=](const LegalityQuery &Query) {
const LLT &DstTy = Query.Types[0];
const LLT &SrcTy = Query.Types[1];

return SrcTy.isVector() && DstTy.isScalar() &&
DstTy == SrcTy.getElementType();
// Handle 2 cases
// 1: vector of size 256 to 2 vectors of size 128
// 2: Unmerge vector to scalar outputs
return (DstTy.isVector() && SrcTy.getSizeInBits() == 256 &&
DstTy.getSizeInBits() == 128) ||
(SrcTy.isVector() && DstTy.isScalar() &&
DstTy == SrcTy.getElementType());
})
.unsupportedIf(IsNotValidDestinationVector)
.legalIf(isValidVectorMergeUnmergeOp(1, 0));
Expand Down
51 changes: 51 additions & 0 deletions llvm/test/CodeGen/AIE/aie2p/GlobalIsel/legalize-unmerge-values.mir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,3 +77,54 @@ body: |
%1:_(s64), %2:_(s64), %3:_(s64), %4:_(s64), %5:_(s64), %6:_(s64), %7:_(s64), %8:_(s64), %9:_(s64), %10:_(s64), %11:_(s64), %12:_(s64), %13:_(s64), %14:_(s64), %15:_(s64), %16:_(s64), %17:_(s64), %18:_(s64), %19:_(s64), %20:_(s64), %21:_(s64), %22:_(s64), %23:_(s64), %24:_(s64), %25:_(s64), %26:_(s64), %27:_(s64), %28:_(s64), %29:_(s64), %30:_(s64), %31:_(s64), %32:_(s64) = G_UNMERGE_VALUES %0:_(s2048)
PseudoRET implicit $lr, implicit %1
...
---
name: test__s128_s256_v8
body: |
bb.0.entry:
; CHECK-LABEL: name: test__s128_s256_v8
; CHECK: [[DEF:%[0-9]+]]:_(<8 x s32>) = G_IMPLICIT_DEF
; CHECK-NEXT: [[AIE_UNPAD_VECTOR:%[0-9]+]]:_(<4 x s32>) = G_AIE_UNPAD_VECTOR [[DEF]](<8 x s32>)
; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 16
; CHECK-NEXT: [[CONCAT_VECTORS:%[0-9]+]]:_(<16 x s32>) = G_CONCAT_VECTORS [[DEF]](<8 x s32>), [[DEF]](<8 x s32>)
; CHECK-NEXT: [[DEF1:%[0-9]+]]:_(<16 x s32>) = G_IMPLICIT_DEF
; CHECK-NEXT: [[AIE_VSHIFT_RIGHT:%[0-9]+]]:_(<16 x s32>) = G_AIE_VSHIFT_RIGHT [[CONCAT_VECTORS]], [[DEF1]], [[C]](s32)
; CHECK-NEXT: [[AIE_UNPAD_VECTOR1:%[0-9]+]]:_(<4 x s32>) = G_AIE_UNPAD_VECTOR [[AIE_VSHIFT_RIGHT]](<16 x s32>)
; CHECK-NEXT: PseudoRET implicit $lr, implicit [[AIE_UNPAD_VECTOR]](<4 x s32>), implicit [[AIE_UNPAD_VECTOR1]](<4 x s32>)
%0:_(<8 x s32>) = G_IMPLICIT_DEF
%1:_(<4 x s32>), %2:_(<4 x s32>) = G_UNMERGE_VALUES %0:_(<8 x s32>)
PseudoRET implicit $lr, implicit %1, implicit %2
...
---
name: test__s128_s256_v16
body: |
bb.0.entry:
; CHECK-LABEL: name: test__s128_s256_v16
; CHECK: [[DEF:%[0-9]+]]:_(<16 x s16>) = G_IMPLICIT_DEF
; CHECK-NEXT: [[AIE_UNPAD_VECTOR:%[0-9]+]]:_(<8 x s16>) = G_AIE_UNPAD_VECTOR [[DEF]](<16 x s16>)
; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 16
; CHECK-NEXT: [[CONCAT_VECTORS:%[0-9]+]]:_(<32 x s16>) = G_CONCAT_VECTORS [[DEF]](<16 x s16>), [[DEF]](<16 x s16>)
; CHECK-NEXT: [[DEF1:%[0-9]+]]:_(<32 x s16>) = G_IMPLICIT_DEF
; CHECK-NEXT: [[AIE_VSHIFT_RIGHT:%[0-9]+]]:_(<32 x s16>) = G_AIE_VSHIFT_RIGHT [[CONCAT_VECTORS]], [[DEF1]], [[C]](s32)
; CHECK-NEXT: [[AIE_UNPAD_VECTOR1:%[0-9]+]]:_(<8 x s16>) = G_AIE_UNPAD_VECTOR [[AIE_VSHIFT_RIGHT]](<32 x s16>)
; CHECK-NEXT: PseudoRET implicit $lr, implicit [[AIE_UNPAD_VECTOR]](<8 x s16>), implicit [[AIE_UNPAD_VECTOR1]](<8 x s16>)
%0:_(<16 x s16>) = G_IMPLICIT_DEF
%1:_(<8 x s16>), %2:_(<8 x s16>) = G_UNMERGE_VALUES %0:_(<16 x s16>)
PseudoRET implicit $lr, implicit %1, implicit %2
...
---
name: test__s128_s256_v32
body: |
bb.0.entry:
; CHECK-LABEL: name: test__s128_s256_v32
; CHECK: [[DEF:%[0-9]+]]:_(<32 x s8>) = G_IMPLICIT_DEF
; CHECK-NEXT: [[AIE_UNPAD_VECTOR:%[0-9]+]]:_(<16 x s8>) = G_AIE_UNPAD_VECTOR [[DEF]](<32 x s8>)
; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 16
; CHECK-NEXT: [[CONCAT_VECTORS:%[0-9]+]]:_(<64 x s8>) = G_CONCAT_VECTORS [[DEF]](<32 x s8>), [[DEF]](<32 x s8>)
; CHECK-NEXT: [[DEF1:%[0-9]+]]:_(<64 x s8>) = G_IMPLICIT_DEF
; CHECK-NEXT: [[AIE_VSHIFT_RIGHT:%[0-9]+]]:_(<64 x s8>) = G_AIE_VSHIFT_RIGHT [[CONCAT_VECTORS]], [[DEF1]], [[C]](s32)
; CHECK-NEXT: [[AIE_UNPAD_VECTOR1:%[0-9]+]]:_(<16 x s8>) = G_AIE_UNPAD_VECTOR [[AIE_VSHIFT_RIGHT]](<64 x s8>)
; CHECK-NEXT: PseudoRET implicit $lr, implicit [[AIE_UNPAD_VECTOR]](<16 x s8>), implicit [[AIE_UNPAD_VECTOR1]](<16 x s8>)
%0:_(<32 x s8>) = G_IMPLICIT_DEF
%1:_(<16 x s8>), %2:_(<16 x s8>) = G_UNMERGE_VALUES %0:_(<32 x s8>)
PseudoRET implicit $lr, implicit %1, implicit %2
...

0 comments on commit d6e46e6

Please sign in to comment.