| //===-- X86LowerAMXIntrinsics.cpp -X86 Scalarize AMX Intrinsics------------===// |
| // |
| // Part of the LLVM Project, 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 |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| /// \file Pass to transform amx intrinsics to scalar operations. |
| /// This pass is always enabled and it skips when it is not -O0 and has no |
| /// optnone attributes. With -O0 or optnone attribute, the def of shape to amx |
| /// intrinsics is near the amx intrinsics code. We are not able to find a |
| /// point which post-dominate all the shape and dominate all amx intrinsics. |
| /// To decouple the dependency of the shape, we transform amx intrinsics |
| /// to scalar operation, so that compiling doesn't fail. In long term, we |
| /// should improve fast register allocation to allocate amx register. |
| //===----------------------------------------------------------------------===// |
| // |
| #include "X86.h" |
| #include "llvm/ADT/DenseSet.h" |
| #include "llvm/ADT/PostOrderIterator.h" |
| #include "llvm/Analysis/DomTreeUpdater.h" |
| #include "llvm/Analysis/OptimizationRemarkEmitter.h" |
| #include "llvm/Analysis/TargetTransformInfo.h" |
| #include "llvm/CodeGen/Passes.h" |
| #include "llvm/CodeGen/TargetPassConfig.h" |
| #include "llvm/CodeGen/ValueTypes.h" |
| #include "llvm/IR/DataLayout.h" |
| #include "llvm/IR/Function.h" |
| #include "llvm/IR/IRBuilder.h" |
| #include "llvm/IR/Instructions.h" |
| #include "llvm/IR/IntrinsicInst.h" |
| #include "llvm/IR/IntrinsicsX86.h" |
| #include "llvm/IR/PatternMatch.h" |
| #include "llvm/InitializePasses.h" |
| #include "llvm/Pass.h" |
| #include "llvm/Support/CommandLine.h" |
| #include "llvm/Target/TargetMachine.h" |
| #include "llvm/Transforms/Utils/BasicBlockUtils.h" |
| #include "llvm/Transforms/Utils/LoopUtils.h" |
| |
| using namespace llvm; |
| using namespace PatternMatch; |
| |
| #define DEBUG_TYPE "lower-amx-intrinsics" |
| |
| #ifndef NDEBUG |
| static bool isV256I32Ty(Type *Ty) { |
| if (auto *FVT = dyn_cast<FixedVectorType>(Ty)) |
| return FVT->getNumElements() == 256 && |
| FVT->getElementType()->isIntegerTy(32); |
| return false; |
| } |
| #endif |
| |
| static cl::opt<bool> |
| X86ScalarizeAMX("enable-x86-scalar-amx", cl::init(false), cl::Hidden, |
| cl::desc("X86: enable AMX scalarizition.")); |
| |
| namespace { |
| class X86LowerAMXIntrinsics { |
| Function &Func; |
| |
| public: |
| X86LowerAMXIntrinsics(Function &F, DomTreeUpdater &DomTU, LoopInfo *LoopI) |
| : Func(F), DTU(DomTU), LI(LoopI) {} |
| bool visit(); |
| |
| private: |
| DomTreeUpdater &DTU; |
| LoopInfo *LI; |
| BasicBlock *createLoop(BasicBlock *Preheader, BasicBlock *Exit, Value *Bound, |
| Value *Step, StringRef Name, IRBuilderBase &B, |
| Loop *L); |
| template <bool IsTileLoad> |
| Value *createTileLoadStoreLoops(BasicBlock *Start, BasicBlock *End, |
| IRBuilderBase &B, Value *Row, Value *Col, |
| Value *Ptr, Value *Stride, Value *Tile); |
| template <Intrinsic::ID IntrID> |
| typename std::enable_if<IntrID == Intrinsic::x86_tdpbssd_internal || |
| IntrID == Intrinsic::x86_tdpbsud_internal || |
| IntrID == Intrinsic::x86_tdpbusd_internal || |
| IntrID == Intrinsic::x86_tdpbuud_internal || |
| IntrID == Intrinsic::x86_tdpbf16ps_internal, |
| Value *>::type |
| createTileDPLoops(BasicBlock *Start, BasicBlock *End, IRBuilderBase &B, |
| Value *Row, Value *Col, Value *K, Value *Acc, Value *LHS, |
| Value *RHS); |
| template <bool IsTileLoad> |
| bool lowerTileLoadStore(Instruction *TileLoadStore); |
| template <Intrinsic::ID IntrID> |
| typename std::enable_if<IntrID == Intrinsic::x86_tdpbssd_internal || |
| IntrID == Intrinsic::x86_tdpbsud_internal || |
| IntrID == Intrinsic::x86_tdpbusd_internal || |
| IntrID == Intrinsic::x86_tdpbuud_internal || |
| IntrID == Intrinsic::x86_tdpbf16ps_internal, |
| bool>::type |
| lowerTileDP(Instruction *TileDP); |
| bool lowerTileZero(Instruction *TileZero); |
| }; |
| } // anonymous namespace |
| |
| BasicBlock *X86LowerAMXIntrinsics::createLoop(BasicBlock *Preheader, |
| BasicBlock *Exit, Value *Bound, |
| Value *Step, StringRef Name, |
| IRBuilderBase &B, Loop *L) { |
| LLVMContext &Ctx = Preheader->getContext(); |
| BasicBlock *Header = |
| BasicBlock::Create(Ctx, Name + ".header", Preheader->getParent(), Exit); |
| BasicBlock *Body = |
| BasicBlock::Create(Ctx, Name + ".body", Header->getParent(), Exit); |
| BasicBlock *Latch = |
| BasicBlock::Create(Ctx, Name + ".latch", Header->getParent(), Exit); |
| |
| Type *I16Ty = Type::getInt16Ty(Ctx); |
| BranchInst::Create(Body, Header); |
| BranchInst::Create(Latch, Body); |
| PHINode *IV = |
| PHINode::Create(I16Ty, 2, Name + ".iv", Header->getTerminator()); |
| IV->addIncoming(ConstantInt::get(I16Ty, 0), Preheader); |
| |
| B.SetInsertPoint(Latch); |
| Value *Inc = B.CreateAdd(IV, Step, Name + ".step"); |
| Value *Cond = B.CreateICmpNE(Inc, Bound, Name + ".cond"); |
| BranchInst::Create(Header, Exit, Cond, Latch); |
| IV->addIncoming(Inc, Latch); |
| |
| BranchInst *PreheaderBr = cast<BranchInst>(Preheader->getTerminator()); |
| BasicBlock *Tmp = PreheaderBr->getSuccessor(0); |
| PreheaderBr->setSuccessor(0, Header); |
| DTU.applyUpdatesPermissive({ |
| {DominatorTree::Delete, Preheader, Tmp}, |
| {DominatorTree::Insert, Header, Body}, |
| {DominatorTree::Insert, Body, Latch}, |
| {DominatorTree::Insert, Latch, Header}, |
| {DominatorTree::Insert, Latch, Exit}, |
| {DominatorTree::Insert, Preheader, Header}, |
| }); |
| if (LI) { |
| L->addBasicBlockToLoop(Header, *LI); |
| L->addBasicBlockToLoop(Body, *LI); |
| L->addBasicBlockToLoop(Latch, *LI); |
| } |
| return Body; |
| } |
| |
| template <bool IsTileLoad> |
| Value *X86LowerAMXIntrinsics::createTileLoadStoreLoops( |
| BasicBlock *Start, BasicBlock *End, IRBuilderBase &B, Value *Row, |
| Value *Col, Value *Ptr, Value *Stride, Value *Tile) { |
| std::string IntrinName = IsTileLoad ? "tileload" : "tilestore"; |
| Loop *RowLoop = nullptr; |
| Loop *ColLoop = nullptr; |
| if (LI) { |
| RowLoop = LI->AllocateLoop(); |
| ColLoop = LI->AllocateLoop(); |
| RowLoop->addChildLoop(ColLoop); |
| if (Loop *ParentL = LI->getLoopFor(Start)) |
| ParentL->addChildLoop(RowLoop); |
| else |
| LI->addTopLevelLoop(RowLoop); |
| } |
| |
| BasicBlock *RowBody = createLoop(Start, End, Row, B.getInt16(1), |
| IntrinName + ".scalarize.rows", B, RowLoop); |
| BasicBlock *RowLatch = RowBody->getSingleSuccessor(); |
| |
| BasicBlock *ColBody = createLoop(RowBody, RowLatch, Col, B.getInt16(1), |
| IntrinName + ".scalarize.cols", B, ColLoop); |
| |
| BasicBlock *ColLoopLatch = ColBody->getSingleSuccessor(); |
| BasicBlock *ColLoopHeader = ColBody->getSinglePredecessor(); |
| BasicBlock *RowLoopHeader = RowBody->getSinglePredecessor(); |
| Value *CurrentRow = &*RowLoopHeader->begin(); |
| Value *CurrentCol = &*ColLoopHeader->begin(); |
| Type *EltTy = B.getInt32Ty(); |
| FixedVectorType *V256I32Ty = FixedVectorType::get(EltTy, 256); |
| |
| // Common part for tileload and tilestore |
| // *.scalarize.cols.body: |
| // Calculate %idxmem and %idxvec |
| B.SetInsertPoint(ColBody->getTerminator()); |
| Value *CurrentRowZExt = B.CreateZExt(CurrentRow, Stride->getType()); |
| Value *CurrentColZExt = B.CreateZExt(CurrentCol, Stride->getType()); |
| Value *Offset = |
| B.CreateAdd(B.CreateMul(CurrentRowZExt, Stride), CurrentColZExt); |
| unsigned AS = cast<PointerType>(Ptr->getType())->getAddressSpace(); |
| Value *EltBasePtr = B.CreatePointerCast(Ptr, PointerType::get(EltTy, AS)); |
| Value *EltPtr = B.CreateGEP(EltTy, EltBasePtr, Offset); |
| Value *Idx = B.CreateAdd(B.CreateMul(CurrentRow, B.getInt16(16)), CurrentCol); |
| if (IsTileLoad) { |
| // tileload.scalarize.rows.header: |
| // %vec.phi.row = phi <256 x i32> [ zeroinitializer, %entry ], [ %ResVec, |
| // %tileload.scalarize.rows.latch ] |
| B.SetInsertPoint(RowLoopHeader->getTerminator()); |
| Value *VecZero = Constant::getNullValue(V256I32Ty); |
| PHINode *VecCPhiRowLoop = B.CreatePHI(V256I32Ty, 2, "vec.phi.row"); |
| VecCPhiRowLoop->addIncoming(VecZero, Start); |
| |
| // tileload.scalarize.cols.header: |
| // %vec.phi = phi <256 x i32> [ %vec.phi.row, %tileload.scalarize.rows.body |
| // ], [ %ResVec, %tileload.scalarize.cols.latch ] |
| B.SetInsertPoint(ColLoopHeader->getTerminator()); |
| PHINode *VecPhi = B.CreatePHI(V256I32Ty, 2, "vec.phi"); |
| VecPhi->addIncoming(VecCPhiRowLoop, RowBody); |
| |
| // tileload.scalarize.cols.body: |
| // Calculate %idxmem and %idxvec |
| // %eltptr = getelementptr i32, i32* %base, i64 %idxmem |
| // %elt = load i32, i32* %ptr |
| // %ResVec = insertelement <256 x i32> %vec.phi, i32 %elt, i16 %idxvec |
| B.SetInsertPoint(ColBody->getTerminator()); |
| Value *Elt = B.CreateLoad(EltTy, EltPtr); |
| Value *ResVec = B.CreateInsertElement(VecPhi, Elt, Idx); |
| VecPhi->addIncoming(ResVec, ColLoopLatch); |
| VecCPhiRowLoop->addIncoming(ResVec, RowLatch); |
| |
| return ResVec; |
| } else { |
| auto *BitCast = cast<BitCastInst>(Tile); |
| Value *Vec = BitCast->getOperand(0); |
| assert(isV256I32Ty(Vec->getType()) && "bitcast from non-v256i32 to x86amx"); |
| // tilestore.scalarize.cols.body: |
| // %mul = mul i16 %row.iv, i16 16 |
| // %idx = add i16 %mul, i16 %col.iv |
| // %vec = extractelement <16 x i32> %vec, i16 %idx |
| // store i32 %vec, i32* %ptr |
| B.SetInsertPoint(ColBody->getTerminator()); |
| Value *Elt = B.CreateExtractElement(Vec, Idx); |
| |
| B.CreateStore(Elt, EltPtr); |
| return nullptr; |
| } |
| } |
| |
| template <Intrinsic::ID IntrID> |
| typename std::enable_if<IntrID == Intrinsic::x86_tdpbssd_internal || |
| IntrID == Intrinsic::x86_tdpbsud_internal || |
| IntrID == Intrinsic::x86_tdpbusd_internal || |
| IntrID == Intrinsic::x86_tdpbuud_internal || |
| IntrID == Intrinsic::x86_tdpbf16ps_internal, |
| Value *>::type |
| X86LowerAMXIntrinsics::createTileDPLoops(BasicBlock *Start, BasicBlock *End, |
| IRBuilderBase &B, Value *Row, |
| Value *Col, Value *K, Value *Acc, |
| Value *LHS, Value *RHS) { |
| std::string IntrinName; |
| switch (IntrID) { |
| case Intrinsic::x86_tdpbssd_internal: |
| IntrinName = "tiledpbssd"; |
| break; |
| case Intrinsic::x86_tdpbsud_internal: |
| IntrinName = "tiledpbsud"; |
| break; |
| case Intrinsic::x86_tdpbusd_internal: |
| IntrinName = "tiledpbusd"; |
| break; |
| case Intrinsic::x86_tdpbuud_internal: |
| IntrinName = "tiledpbuud"; |
| break; |
| case Intrinsic::x86_tdpbf16ps_internal: |
| IntrinName = "tiledpbf16ps"; |
| break; |
| } |
| Loop *RowLoop = nullptr; |
| Loop *ColLoop = nullptr; |
| Loop *InnerLoop = nullptr; |
| if (LI) { |
| RowLoop = LI->AllocateLoop(); |
| ColLoop = LI->AllocateLoop(); |
| InnerLoop = LI->AllocateLoop(); |
| ColLoop->addChildLoop(InnerLoop); |
| RowLoop->addChildLoop(ColLoop); |
| if (Loop *ParentL = LI->getLoopFor(Start)) |
| ParentL->addChildLoop(RowLoop); |
| else |
| LI->addTopLevelLoop(RowLoop); |
| } |
| |
| BasicBlock *RowBody = createLoop(Start, End, Row, B.getInt16(1), |
| IntrinName + ".scalarize.rows", B, RowLoop); |
| BasicBlock *RowLatch = RowBody->getSingleSuccessor(); |
| |
| BasicBlock *ColBody = createLoop(RowBody, RowLatch, Col, B.getInt16(1), |
| IntrinName + ".scalarize.cols", B, ColLoop); |
| |
| BasicBlock *ColLoopLatch = ColBody->getSingleSuccessor(); |
| |
| B.SetInsertPoint(ColBody->getTerminator()); |
| BasicBlock *InnerBody = |
| createLoop(ColBody, ColLoopLatch, K, B.getInt16(1), |
| IntrinName + ".scalarize.inner", B, InnerLoop); |
| |
| BasicBlock *ColLoopHeader = ColBody->getSinglePredecessor(); |
| BasicBlock *RowLoopHeader = RowBody->getSinglePredecessor(); |
| BasicBlock *InnerLoopHeader = InnerBody->getSinglePredecessor(); |
| BasicBlock *InnerLoopLatch = InnerBody->getSingleSuccessor(); |
| Value *CurrentRow = &*RowLoopHeader->begin(); |
| Value *CurrentCol = &*ColLoopHeader->begin(); |
| Value *CurrentInner = &*InnerLoopHeader->begin(); |
| |
| FixedVectorType *V256I32Ty = FixedVectorType::get(B.getInt32Ty(), 256); |
| auto *BitCastAcc = cast<BitCastInst>(Acc); |
| Value *VecC = BitCastAcc->getOperand(0); |
| assert(isV256I32Ty(VecC->getType()) && "bitcast from non-v256i32 to x86amx"); |
| // TODO else create BitCast from x86amx to v256i32. |
| // Store x86amx to memory, and reload from memory |
| // to vector. However with -O0, it doesn't happen. |
| auto *BitCastLHS = cast<BitCastInst>(LHS); |
| Value *VecA = BitCastLHS->getOperand(0); |
| assert(isV256I32Ty(VecA->getType()) && "bitcast from non-v256i32 to x86amx"); |
| auto *BitCastRHS = cast<BitCastInst>(RHS); |
| Value *VecB = BitCastRHS->getOperand(0); |
| assert(isV256I32Ty(VecB->getType()) && "bitcast from non-v256i32 to x86amx"); |
| |
| // tiledpbssd.scalarize.rows.header: |
| // %vec.c.phi.row = phi <256 x i32> [ %VecC, %continue ], [ %NewVecC, |
| // %tiledpbssd.scalarize.rows.latch ] |
| |
| // %vec.d.phi.row = phi <256 x i32> [ zeroinitializer, %continue ], [ |
| // %NewVecD, %tiledpbssd.scalarize.rows.latch ] |
| B.SetInsertPoint(RowLoopHeader->getTerminator()); |
| PHINode *VecCPhiRowLoop = B.CreatePHI(V256I32Ty, 2, "vec.c.phi.row"); |
| VecCPhiRowLoop->addIncoming(VecC, Start); |
| Value *VecZero = Constant::getNullValue(V256I32Ty); |
| PHINode *VecDPhiRowLoop = B.CreatePHI(V256I32Ty, 2, "vec.d.phi.row"); |
| VecDPhiRowLoop->addIncoming(VecZero, Start); |
| |
| // tiledpbssd.scalarize.cols.header: |
| // %vec.c.phi.col = phi <256 x i32> [ %vec.c.phi.row, |
| // %tiledpbssd.scalarize.rows.body ], [ %NewVecC, |
| // %tiledpbssd.scalarize.cols.latch ] |
| |
| // %vec.d.phi.col = phi <256 x i32> [ |
| // %vec.d.phi.row, %tiledpbssd.scalarize.rows.body ], [ %NewVecD, |
| // %tiledpbssd.scalarize.cols.latch ] |
| |
| // calculate idxc. |
| B.SetInsertPoint(ColLoopHeader->getTerminator()); |
| PHINode *VecCPhiColLoop = B.CreatePHI(V256I32Ty, 2, "vec.c.phi.col"); |
| VecCPhiColLoop->addIncoming(VecCPhiRowLoop, RowBody); |
| PHINode *VecDPhiColLoop = B.CreatePHI(V256I32Ty, 2, "vec.d.phi.col"); |
| VecDPhiColLoop->addIncoming(VecDPhiRowLoop, RowBody); |
| Value *IdxC = |
| B.CreateAdd(B.CreateMul(CurrentRow, B.getInt16(16)), CurrentCol); |
| |
| // tiledpbssd.scalarize.inner.header: |
| // %vec.c.inner.phi = phi <256 x i32> [ %vec.c.phi.col, |
| // %tiledpbssd.scalarize.cols.body ], [ %NewVecC, |
| // %tiledpbssd.scalarize.inner.latch ] |
| |
| B.SetInsertPoint(InnerLoopHeader->getTerminator()); |
| PHINode *VecCPhi = B.CreatePHI(V256I32Ty, 2, "vec.c.inner.phi"); |
| VecCPhi->addIncoming(VecCPhiColLoop, ColBody); |
| |
| B.SetInsertPoint(InnerBody->getTerminator()); |
| Value *IdxA = |
| B.CreateAdd(B.CreateMul(CurrentRow, B.getInt16(16)), CurrentInner); |
| Value *IdxB = |
| B.CreateAdd(B.CreateMul(CurrentInner, B.getInt16(16)), CurrentCol); |
| Value *NewVecC = nullptr; |
| |
| if (IntrID != Intrinsic::x86_tdpbf16ps_internal) { |
| // tiledpbssd.scalarize.inner.body: |
| // calculate idxa, idxb |
| // %eltc = extractelement <256 x i32> %vec.c.inner.phi, i16 %idxc |
| // %elta = extractelement <256 x i32> %veca, i16 %idxa |
| // %eltav4i8 = bitcast i32 %elta to <4 x i8> |
| // %eltb = extractelement <256 x i32> %vecb, i16 %idxb |
| // %eltbv4i8 = bitcast i32 %eltb to <4 x i8> |
| // %eltav4i32 = sext <4 x i8> %eltav4i8 to <4 x i32> |
| // %eltbv4i32 = sext <4 x i8> %eltbv4i8 to <4 x i32> |
| // %mulab = mul <4 x i32> %eltbv4i32, %eltav4i32 |
| // %acc = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> %131) |
| // %neweltc = add i32 %elt, %acc |
| // %NewVecC = insertelement <256 x i32> %vec.c.inner.phi, i32 %neweltc, |
| // i16 %idxc |
| FixedVectorType *V4I8Ty = FixedVectorType::get(B.getInt8Ty(), 4); |
| FixedVectorType *V4I32Ty = FixedVectorType::get(B.getInt32Ty(), 4); |
| Value *EltC = B.CreateExtractElement(VecCPhi, IdxC); |
| Value *EltA = B.CreateExtractElement(VecA, IdxA); |
| Value *SubVecA = B.CreateBitCast(EltA, V4I8Ty); |
| Value *EltB = B.CreateExtractElement(VecB, IdxB); |
| Value *SubVecB = B.CreateBitCast(EltB, V4I8Ty); |
| Value *SEXTSubVecB = nullptr; |
| Value *SEXTSubVecA = nullptr; |
| switch (IntrID) { |
| case Intrinsic::x86_tdpbssd_internal: |
| SEXTSubVecB = B.CreateSExt(SubVecB, V4I32Ty); |
| SEXTSubVecA = B.CreateSExt(SubVecA, V4I32Ty); |
| break; |
| case Intrinsic::x86_tdpbsud_internal: |
| SEXTSubVecB = B.CreateZExt(SubVecB, V4I32Ty); |
| SEXTSubVecA = B.CreateSExt(SubVecA, V4I32Ty); |
| break; |
| case Intrinsic::x86_tdpbusd_internal: |
| SEXTSubVecB = B.CreateSExt(SubVecB, V4I32Ty); |
| SEXTSubVecA = B.CreateZExt(SubVecA, V4I32Ty); |
| break; |
| case Intrinsic::x86_tdpbuud_internal: |
| SEXTSubVecB = B.CreateZExt(SubVecB, V4I32Ty); |
| SEXTSubVecA = B.CreateZExt(SubVecA, V4I32Ty); |
| break; |
| default: |
| llvm_unreachable("Invalid intrinsic ID!"); |
| } |
| Value *SubVecR = B.CreateAddReduce(B.CreateMul(SEXTSubVecA, SEXTSubVecB)); |
| Value *ResElt = B.CreateAdd(EltC, SubVecR); |
| NewVecC = B.CreateInsertElement(VecCPhi, ResElt, IdxC); |
| } else { |
| // tiledpbf16ps.scalarize.inner.body: |
| // calculate idxa, idxb, idxc |
| // %eltc = extractelement <256 x i32> %vec.c.inner.phi, i16 %idxc |
| // %eltcf32 = bitcast i32 %eltc to float |
| // %elta = extractelement <256 x i32> %veca, i16 %idxa |
| // %eltav2i16 = bitcast i32 %elta to <2 x i16> |
| // %eltb = extractelement <256 x i32> %vecb, i16 %idxb |
| // %eltbv2i16 = bitcast i32 %eltb to <2 x i16> |
| // %shufflea = shufflevector <2 x i16> %elta, <2 x i16> zeroinitializer, <4 |
| // x i32> <i32 2, i32 0, i32 3, i32 1> |
| // %eltav2f32 = bitcast <4 x i16> %shufflea to <2 x float> |
| // %shuffleb = shufflevector <2 x i16> %eltb, <2 xi16> zeroinitializer, <4 x |
| // i32> <i32 2, i32 0, i32 3, i32 1> |
| // %eltbv2f32 = bitcast <4 x i16> %shuffleb to <2 x float> |
| // %mulab = fmul <2 x float> %eltav2f32, %eltbv2f32 |
| // %acc = call float |
| // @llvm.vector.reduce.fadd.v2f32(float %eltcf32, <2 x float> %mulab) |
| // %neweltc = bitcast float %acc to i32 |
| // %NewVecC = insertelement <256 x i32> %vec.c.inner.phi, i32 %neweltc, |
| // i16 %idxc |
| // %NewVecD = insertelement <256 x i32> %vec.d.inner.phi, i32 %neweltc, |
| // i16 %idxc |
| FixedVectorType *V2I16Ty = FixedVectorType::get(B.getInt16Ty(), 2); |
| FixedVectorType *V2F32Ty = FixedVectorType::get(B.getFloatTy(), 2); |
| Value *EltC = B.CreateExtractElement(VecCPhi, IdxC); |
| Value *EltCF32 = B.CreateBitCast(EltC, B.getFloatTy()); |
| Value *EltA = B.CreateExtractElement(VecA, IdxA); |
| Value *SubVecA = B.CreateBitCast(EltA, V2I16Ty); |
| Value *EltB = B.CreateExtractElement(VecB, IdxB); |
| Value *SubVecB = B.CreateBitCast(EltB, V2I16Ty); |
| Value *ZeroV2I16 = Constant::getNullValue(V2I16Ty); |
| int ShuffleMask[4] = {2, 0, 3, 1}; |
| auto ShuffleArray = makeArrayRef(ShuffleMask); |
| Value *AV2F32 = B.CreateBitCast( |
| B.CreateShuffleVector(SubVecA, ZeroV2I16, ShuffleArray), V2F32Ty); |
| Value *BV2F32 = B.CreateBitCast( |
| B.CreateShuffleVector(SubVecB, ZeroV2I16, ShuffleArray), V2F32Ty); |
| Value *SubVecR = B.CreateFAddReduce(EltCF32, B.CreateFMul(AV2F32, BV2F32)); |
| Value *ResElt = B.CreateBitCast(SubVecR, B.getInt32Ty()); |
| NewVecC = B.CreateInsertElement(VecCPhi, ResElt, IdxC); |
| } |
| |
| // tiledpbssd.scalarize.cols.latch: |
| // %NewEltC = extractelement <256 x i32> %vec.c.phi.col, i16 %idxc |
| // %NewVecD = insertelement <256 x i32> %vec.d.phi.col, i32 %NewEltC, |
| // i16 %idxc |
| B.SetInsertPoint(ColLoopLatch->getTerminator()); |
| Value *NewEltC = B.CreateExtractElement(NewVecC, IdxC); |
| Value *NewVecD = B.CreateInsertElement(VecDPhiColLoop, NewEltC, IdxC); |
| |
| VecCPhi->addIncoming(NewVecC, InnerLoopLatch); |
| VecCPhiRowLoop->addIncoming(NewVecC, RowLatch); |
| VecCPhiColLoop->addIncoming(NewVecC, ColLoopLatch); |
| VecDPhiRowLoop->addIncoming(NewVecD, RowLatch); |
| VecDPhiColLoop->addIncoming(NewVecD, ColLoopLatch); |
| |
| return NewVecD; |
| } |
| |
| template <Intrinsic::ID IntrID> |
| typename std::enable_if<IntrID == Intrinsic::x86_tdpbssd_internal || |
| IntrID == Intrinsic::x86_tdpbsud_internal || |
| IntrID == Intrinsic::x86_tdpbusd_internal || |
| IntrID == Intrinsic::x86_tdpbuud_internal || |
| IntrID == Intrinsic::x86_tdpbf16ps_internal, |
| bool>::type |
| X86LowerAMXIntrinsics::lowerTileDP(Instruction *TileDP) { |
| Value *M, *N, *K, *C, *A, *B; |
| match(TileDP, m_Intrinsic<IntrID>(m_Value(M), m_Value(N), m_Value(K), |
| m_Value(C), m_Value(A), m_Value(B))); |
| Instruction *InsertI = TileDP; |
| IRBuilder<> PreBuilder(TileDP); |
| PreBuilder.SetInsertPoint(TileDP); |
| // We visit the loop with (m, n/4, k/4): |
| // %n_dword = lshr i16 %n, 2 |
| // %k_dword = lshr i16 %k, 2 |
| Value *NDWord = PreBuilder.CreateLShr(N, PreBuilder.getInt16(2)); |
| Value *KDWord = PreBuilder.CreateLShr(K, PreBuilder.getInt16(2)); |
| BasicBlock *Start = InsertI->getParent(); |
| BasicBlock *End = |
| SplitBlock(InsertI->getParent(), InsertI, &DTU, LI, nullptr, "continue"); |
| IRBuilder<> Builder(TileDP); |
| Value *ResVec = createTileDPLoops<IntrID>(Start, End, Builder, M, NDWord, |
| KDWord, C, A, B); |
| // we cannot assume there always be bitcast after tiledpbssd. So we need to |
| // insert one bitcast as required |
| Builder.SetInsertPoint(End->getFirstNonPHI()); |
| Value *ResAMX = |
| Builder.CreateBitCast(ResVec, Type::getX86_AMXTy(Builder.getContext())); |
| // Delete TileDP intrinsic and do some clean-up. |
| for (Use &U : llvm::make_early_inc_range(TileDP->uses())) { |
| Instruction *I = cast<Instruction>(U.getUser()); |
| Value *Vec; |
| if (match(I, m_BitCast(m_Value(Vec)))) { |
| I->replaceAllUsesWith(ResVec); |
| I->eraseFromParent(); |
| } |
| } |
| TileDP->replaceAllUsesWith(ResAMX); |
| TileDP->eraseFromParent(); |
| return true; |
| } |
| |
| template <bool IsTileLoad> |
| bool X86LowerAMXIntrinsics::lowerTileLoadStore(Instruction *TileLoadStore) { |
| Value *M, *N, *Ptr, *Stride, *Tile; |
| if (IsTileLoad) |
| match(TileLoadStore, |
| m_Intrinsic<Intrinsic::x86_tileloadd64_internal>( |
| m_Value(M), m_Value(N), m_Value(Ptr), m_Value(Stride))); |
| else |
| match(TileLoadStore, m_Intrinsic<Intrinsic::x86_tilestored64_internal>( |
| m_Value(M), m_Value(N), m_Value(Ptr), |
| m_Value(Stride), m_Value(Tile))); |
| |
| Instruction *InsertI = TileLoadStore; |
| IRBuilder<> PreBuilder(TileLoadStore); |
| PreBuilder.SetInsertPoint(TileLoadStore); |
| Value *NDWord = PreBuilder.CreateLShr(N, PreBuilder.getInt16(2)); |
| Value *StrideDWord = PreBuilder.CreateLShr(Stride, PreBuilder.getInt64(2)); |
| BasicBlock *Start = InsertI->getParent(); |
| BasicBlock *End = |
| SplitBlock(InsertI->getParent(), InsertI, &DTU, LI, nullptr, "continue"); |
| IRBuilder<> Builder(TileLoadStore); |
| Value *ResVec = createTileLoadStoreLoops<IsTileLoad>( |
| Start, End, Builder, M, NDWord, Ptr, StrideDWord, |
| IsTileLoad ? nullptr : Tile); |
| if (IsTileLoad) { |
| // we cannot assume there always be bitcast after tileload. So we need to |
| // insert one bitcast as required |
| Builder.SetInsertPoint(End->getFirstNonPHI()); |
| Value *ResAMX = |
| Builder.CreateBitCast(ResVec, Type::getX86_AMXTy(Builder.getContext())); |
| // Delete tileloadd6 intrinsic and do some clean-up |
| for (Use &U : llvm::make_early_inc_range(TileLoadStore->uses())) { |
| Instruction *I = cast<Instruction>(U.getUser()); |
| Value *Vec; |
| if (match(I, m_BitCast(m_Value(Vec)))) { |
| I->replaceAllUsesWith(ResVec); |
| I->eraseFromParent(); |
| } |
| } |
| TileLoadStore->replaceAllUsesWith(ResAMX); |
| } |
| TileLoadStore->eraseFromParent(); |
| return true; |
| } |
| |
| bool X86LowerAMXIntrinsics::lowerTileZero(Instruction *TileZero) { |
| IRBuilder<> Builder(TileZero); |
| FixedVectorType *V256I32Ty = FixedVectorType::get(Builder.getInt32Ty(), 256); |
| Value *VecZero = Constant::getNullValue(V256I32Ty); |
| for (Use &U : llvm::make_early_inc_range(TileZero->uses())) { |
| Instruction *I = cast<Instruction>(U.getUser()); |
| Value *Vec; |
| if (match(I, m_BitCast(m_Value(Vec)))) { |
| I->replaceAllUsesWith(VecZero); |
| I->eraseFromParent(); |
| } |
| } |
| TileZero->eraseFromParent(); |
| return true; |
| } |
| |
| bool X86LowerAMXIntrinsics::visit() { |
| bool C = false; |
| SmallVector<IntrinsicInst *, 8> WorkList; |
| for (BasicBlock *BB : depth_first(&Func)) { |
| for (BasicBlock::iterator II = BB->begin(), IE = BB->end(); II != IE;) { |
| if (auto *Inst = dyn_cast<IntrinsicInst>(&*II++)) { |
| switch (Inst->getIntrinsicID()) { |
| case Intrinsic::x86_tdpbssd_internal: |
| case Intrinsic::x86_tdpbsud_internal: |
| case Intrinsic::x86_tdpbusd_internal: |
| case Intrinsic::x86_tdpbuud_internal: |
| case Intrinsic::x86_tileloadd64_internal: |
| case Intrinsic::x86_tilestored64_internal: |
| case Intrinsic::x86_tilezero_internal: |
| case Intrinsic::x86_tdpbf16ps_internal: |
| WorkList.push_back(Inst); |
| break; |
| default: |
| break; |
| } |
| } |
| } |
| } |
| |
| for (auto *Inst : WorkList) { |
| switch (Inst->getIntrinsicID()) { |
| case Intrinsic::x86_tdpbssd_internal: |
| C = lowerTileDP<Intrinsic::x86_tdpbssd_internal>(Inst) || C; |
| break; |
| case Intrinsic::x86_tdpbsud_internal: |
| C = lowerTileDP<Intrinsic::x86_tdpbsud_internal>(Inst) || C; |
| break; |
| case Intrinsic::x86_tdpbusd_internal: |
| C = lowerTileDP<Intrinsic::x86_tdpbusd_internal>(Inst) || C; |
| break; |
| case Intrinsic::x86_tdpbuud_internal: |
| C = lowerTileDP<Intrinsic::x86_tdpbuud_internal>(Inst) || C; |
| break; |
| case Intrinsic::x86_tdpbf16ps_internal: |
| C = lowerTileDP<Intrinsic::x86_tdpbf16ps_internal>(Inst) || C; |
| break; |
| case Intrinsic::x86_tileloadd64_internal: |
| C = lowerTileLoadStore<true>(Inst) || C; |
| break; |
| case Intrinsic::x86_tilestored64_internal: |
| C = lowerTileLoadStore<false>(Inst) || C; |
| break; |
| case Intrinsic::x86_tilezero_internal: |
| C = lowerTileZero(Inst) || C; |
| break; |
| default: |
| llvm_unreachable("invalid amx intrinsics!"); |
| } |
| } |
| |
| return C; |
| } |
| |
| namespace { |
| class X86LowerAMXIntrinsicsLegacyPass : public FunctionPass { |
| public: |
| static char ID; |
| |
| X86LowerAMXIntrinsicsLegacyPass() : FunctionPass(ID) { |
| initializeX86LowerAMXIntrinsicsLegacyPassPass( |
| *PassRegistry::getPassRegistry()); |
| } |
| |
| bool runOnFunction(Function &F) override { |
| if (!X86ScalarizeAMX) |
| return false; |
| TargetMachine *TM = &getAnalysis<TargetPassConfig>().getTM<TargetMachine>(); |
| if (!F.hasFnAttribute(Attribute::OptimizeNone) && |
| TM->getOptLevel() != CodeGenOpt::None) |
| return false; |
| |
| auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>(); |
| auto *DT = DTWP ? &DTWP->getDomTree() : nullptr; |
| auto *LIWP = getAnalysisIfAvailable<LoopInfoWrapperPass>(); |
| auto *LI = LIWP ? &LIWP->getLoopInfo() : nullptr; |
| DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy); |
| |
| X86LowerAMXIntrinsics LAT(F, DTU, LI); |
| return LAT.visit(); |
| } |
| StringRef getPassName() const override { return "Lower AMX intrinsics"; } |
| |
| void getAnalysisUsage(AnalysisUsage &AU) const override { |
| AU.addPreserved<DominatorTreeWrapperPass>(); |
| AU.addPreserved<LoopInfoWrapperPass>(); |
| AU.addRequired<TargetPassConfig>(); |
| } |
| }; |
| } // namespace |
| |
| static const char PassName[] = "Lower AMX intrinsics"; |
| char X86LowerAMXIntrinsicsLegacyPass::ID = 0; |
| INITIALIZE_PASS_BEGIN(X86LowerAMXIntrinsicsLegacyPass, DEBUG_TYPE, PassName, |
| false, false) |
| INITIALIZE_PASS_DEPENDENCY(TargetPassConfig) |
| INITIALIZE_PASS_END(X86LowerAMXIntrinsicsLegacyPass, DEBUG_TYPE, PassName, |
| false, false) |
| |
| FunctionPass *llvm::createX86LowerAMXIntrinsicsPass() { |
| return new X86LowerAMXIntrinsicsLegacyPass(); |
| } |
