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 20, 2025
1 parent 1c3997e commit a2d74e1
Show file tree
Hide file tree
Showing 4 changed files with 124 additions and 3 deletions.
56 changes: 56 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,57 @@ bool AIELegalizerHelper::legalizeG_BUILD_VECTOR(LegalizerHelper &Helper,
return true;
}

bool AIELegalizerHelper::legalize_unmerge_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 Dst1Reg = MI.getOperand(0).getReg();
const Register Dst2Reg = 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, {Dst1Reg}, {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.buildInstr(TargetOpcode::G_IMPLICIT_DEF, {SrcTy}, {})
.getReg(0);

const LLT Vec512 =
LLT::fixed_vector(SrcTy.getNumElements() * 2, SrcTy.getElementType());

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

// The second input will be ignored. Just create a dummy input
auto ImplicitDef512 =
MIRBuilder.buildInstr(TargetOpcode::G_IMPLICIT_DEF, {Vec512}, {});

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

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

MI.eraseFromParent();
return true;
}

bool AIELegalizerHelper::legalizeG_UNMERGE_VALUES(LegalizerHelper &Helper,
MachineInstr &MI) const {
MachineIRBuilder &MIRBuilder = Helper.MIRBuilder;
Expand All @@ -241,6 +292,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.getSizeInBits() == 128 &&
LastTy.getSizeInBits() == 256)
return legalize_unmerge_128bit(Helper, MI);

assert(LastTy.isVector() &&
(FirstTy.getScalarSizeInBits() * (MI.getNumOperands() - 1)) ==
LastTy.getSizeInBits() &&
Expand Down
1 change: 1 addition & 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,7 @@ class AIELegalizerHelper {
Register SourceReg) const;
bool unpack32BitVector(LegalizerHelper &Helper, MachineInstr &MI,
Register SourceReg) const;
bool legalize_unmerge_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 @@ -550,9 +550,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
60 changes: 60 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
@@ -0,0 +1,60 @@
# 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=legalizer %s -verify-machineinstrs -o - | FileCheck %s

---
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:%[0-9]+]]:_(<16 x s32>) = G_AIE_VSHIFT [[CONCAT_VECTORS]], [[DEF1]], [[C]](s32)
; CHECK-NEXT: [[AIE_UNPAD_VECTOR1:%[0-9]+]]:_(<4 x s32>) = G_AIE_UNPAD_VECTOR [[AIE_VSHIFT]](<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:%[0-9]+]]:_(<32 x s16>) = G_AIE_VSHIFT [[CONCAT_VECTORS]], [[DEF1]], [[C]](s32)
; CHECK-NEXT: [[AIE_UNPAD_VECTOR1:%[0-9]+]]:_(<8 x s16>) = G_AIE_UNPAD_VECTOR [[AIE_VSHIFT]](<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:%[0-9]+]]:_(<64 x s8>) = G_AIE_VSHIFT [[CONCAT_VECTORS]], [[DEF1]], [[C]](s32)
; CHECK-NEXT: [[AIE_UNPAD_VECTOR1:%[0-9]+]]:_(<16 x s8>) = G_AIE_UNPAD_VECTOR [[AIE_VSHIFT]](<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 a2d74e1

Please sign in to comment.