LLVM 19.0.0git
X86LowerAMXIntrinsics.cpp
Go to the documentation of this file.
1//===-- X86LowerAMXIntrinsics.cpp -X86 Scalarize AMX Intrinsics------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9/// \file Pass to transform amx intrinsics to scalar operations.
10/// This pass is always enabled and it skips when it is not -O0 and has no
11/// optnone attributes. With -O0 or optnone attribute, the def of shape to amx
12/// intrinsics is near the amx intrinsics code. We are not able to find a
13/// point which post-dominate all the shape and dominate all amx intrinsics.
14/// To decouple the dependency of the shape, we transform amx intrinsics
15/// to scalar operation, so that compiling doesn't fail. In long term, we
16/// should improve fast register allocation to allocate amx register.
17//===----------------------------------------------------------------------===//
18//
19#include "X86.h"
20#include "llvm/ADT/PostOrderIterator.h"
21#include "llvm/Analysis/DomTreeUpdater.h"
22#include "llvm/Analysis/LoopInfo.h"
23#include "llvm/Analysis/OptimizationRemarkEmitter.h"
24#include "llvm/Analysis/TargetTransformInfo.h"
25#include "llvm/CodeGen/Passes.h"
26#include "llvm/CodeGen/TargetPassConfig.h"
27#include "llvm/CodeGen/ValueTypes.h"
28#include "llvm/IR/DataLayout.h"
29#include "llvm/IR/Function.h"
30#include "llvm/IR/IRBuilder.h"
31#include "llvm/IR/Instructions.h"
32#include "llvm/IR/IntrinsicInst.h"
33#include "llvm/IR/IntrinsicsX86.h"
34#include "llvm/IR/PatternMatch.h"
35#include "llvm/InitializePasses.h"
36#include "llvm/Pass.h"
37#include "llvm/Support/CommandLine.h"
38#include "llvm/Target/TargetMachine.h"
39#include "llvm/Transforms/Utils/BasicBlockUtils.h"
40#include "llvm/Transforms/Utils/LoopUtils.h"
41
42using namespace llvm;
43using namespace PatternMatch;
44
45#define DEBUG_TYPE "lower-amx-intrinsics"
46
47#ifndef NDEBUG
48static bool isV256I32Ty(Type *Ty) {
49 if (auto *FVT = dyn_cast<FixedVectorType>(Ty))
50 return FVT->getNumElements() == 256 &&
51 FVT->getElementType()->isIntegerTy(32);
52 return false;
53}
54#endif
55
56static cl::opt<bool>
57 X86ScalarizeAMX("enable-x86-scalar-amx", cl::init(false), cl::Hidden,
58 cl::desc("X86: enable AMX scalarizition."));
59
60namespace {
61class X86LowerAMXIntrinsics {
62 Function &Func;
63
64public:
65 X86LowerAMXIntrinsics(Function &F, DomTreeUpdater &DomTU, LoopInfo *LoopI)
66 : Func(F), DTU(DomTU), LI(LoopI) {}
67 bool visit();
68
69private:
70 DomTreeUpdater &DTU;
71 LoopInfo *LI;
72 BasicBlock *createLoop(BasicBlock *Preheader, BasicBlock *Exit, Value *Bound,
73 Value *Step, StringRef Name, IRBuilderBase &B,
74 Loop *L);
75 template <bool IsTileLoad>
76 Value *createTileLoadStoreLoops(BasicBlock *Start, BasicBlock *End,
77 IRBuilderBase &B, Value *Row, Value *Col,
78 Value *Ptr, Value *Stride, Value *Tile);
79 template <Intrinsic::ID IntrID>
80 std::enable_if_t<IntrID == Intrinsic::x86_tdpbssd_internal ||
81 IntrID == Intrinsic::x86_tdpbsud_internal ||
82 IntrID == Intrinsic::x86_tdpbusd_internal ||
83 IntrID == Intrinsic::x86_tdpbuud_internal ||
84 IntrID == Intrinsic::x86_tdpbf16ps_internal,
85 Value *>
86 createTileDPLoops(BasicBlock *Start, BasicBlock *End, IRBuilderBase &B,
87 Value *Row, Value *Col, Value *K, Value *Acc, Value *LHS,
88 Value *RHS);
89 template <bool IsTileLoad>
90 bool lowerTileLoadStore(Instruction *TileLoadStore);
91 template <Intrinsic::ID IntrID>
92 std::enable_if_t<IntrID == Intrinsic::x86_tdpbssd_internal ||
93 IntrID == Intrinsic::x86_tdpbsud_internal ||
94 IntrID == Intrinsic::x86_tdpbusd_internal ||
95 IntrID == Intrinsic::x86_tdpbuud_internal ||
96 IntrID == Intrinsic::x86_tdpbf16ps_internal,
97 bool>
98 lowerTileDP(Instruction *TileDP);
99 bool lowerTileZero(Instruction *TileZero);
100};
101} // anonymous namespace
102
103BasicBlock *X86LowerAMXIntrinsics::createLoop(BasicBlock *Preheader,
104 BasicBlock *Exit, Value *Bound,
105 Value *Step, StringRef Name,
106 IRBuilderBase &B, Loop *L) {
107 LLVMContext &Ctx = Preheader->getContext();
108 BasicBlock *Header =
109 BasicBlock::Create(Ctx, Name + ".header", Preheader->getParent(), Exit);
110 BasicBlock *Body =
111 BasicBlock::Create(Ctx, Name + ".body", Header->getParent(), Exit);
112 BasicBlock *Latch =
113 BasicBlock::Create(Ctx, Name + ".latch", Header->getParent(), Exit);
114
115 Type *I16Ty = Type::getInt16Ty(Ctx);
116 BranchInst::Create(Body, Header);
117 BranchInst::Create(Latch, Body);
118 PHINode *IV =
119 PHINode::Create(I16Ty, 2, Name + ".iv", Header->getTerminator()->getIterator());
120 IV->addIncoming(ConstantInt::get(I16Ty, 0), Preheader);
121
122 B.SetInsertPoint(Latch);
123 Value *Inc = B.CreateAdd(IV, Step, Name + ".step");
124 Value *Cond = B.CreateICmpNE(Inc, Bound, Name + ".cond");
125 BranchInst::Create(Header, Exit, Cond, Latch);
126 IV->addIncoming(Inc, Latch);
127
128 BranchInst *PreheaderBr = cast<BranchInst>(Preheader->getTerminator());
129 BasicBlock *Tmp = PreheaderBr->getSuccessor(0);
130 PreheaderBr->setSuccessor(0, Header);
131 DTU.applyUpdatesPermissive({
132 {DominatorTree::Delete, Preheader, Tmp},
133 {DominatorTree::Insert, Header, Body},
134 {DominatorTree::Insert, Body, Latch},
135 {DominatorTree::Insert, Latch, Header},
136 {DominatorTree::Insert, Latch, Exit},
137 {DominatorTree::Insert, Preheader, Header},
138 });
139 if (LI) {
140 L->addBasicBlockToLoop(Header, *LI);
141 L->addBasicBlockToLoop(Body, *LI);
142 L->addBasicBlockToLoop(Latch, *LI);
143 }
144 return Body;
145}
146
147template <bool IsTileLoad>
148Value *X86LowerAMXIntrinsics::createTileLoadStoreLoops(
149 BasicBlock *Start, BasicBlock *End, IRBuilderBase &B, Value *Row,
150 Value *Col, Value *Ptr, Value *Stride, Value *Tile) {
151 std::string IntrinName = IsTileLoad ? "tileload" : "tilestore";
152 Loop *RowLoop = nullptr;
153 Loop *ColLoop = nullptr;
154 if (LI) {
155 RowLoop = LI->AllocateLoop();
156 ColLoop = LI->AllocateLoop();
157 RowLoop->addChildLoop(ColLoop);
158 if (Loop *ParentL = LI->getLoopFor(Start))
159 ParentL->addChildLoop(RowLoop);
160 else
161 LI->addTopLevelLoop(RowLoop);
162 }
163
164 BasicBlock *RowBody = createLoop(Start, End, Row, B.getInt16(1),
165 IntrinName + ".scalarize.rows", B, RowLoop);
166 BasicBlock *RowLatch = RowBody->getSingleSuccessor();
167
168 BasicBlock *ColBody = createLoop(RowBody, RowLatch, Col, B.getInt16(1),
169 IntrinName + ".scalarize.cols", B, ColLoop);
170
171 BasicBlock *ColLoopLatch = ColBody->getSingleSuccessor();
172 BasicBlock *ColLoopHeader = ColBody->getSinglePredecessor();
173 BasicBlock *RowLoopHeader = RowBody->getSinglePredecessor();
174 Value *CurrentRow = &*RowLoopHeader->begin();
175 Value *CurrentCol = &*ColLoopHeader->begin();
176 Type *EltTy = B.getInt32Ty();
177 FixedVectorType *V256I32Ty = FixedVectorType::get(EltTy, 256);
178
179 // Common part for tileload and tilestore
180 // *.scalarize.cols.body:
181 // Calculate %idxmem and %idxvec
182 B.SetInsertPoint(ColBody->getTerminator());
183 Value *CurrentRowZExt = B.CreateZExt(CurrentRow, Stride->getType());
184 Value *CurrentColZExt = B.CreateZExt(CurrentCol, Stride->getType());
185 Value *Offset =
186 B.CreateAdd(B.CreateMul(CurrentRowZExt, Stride), CurrentColZExt);
187 Value *EltPtr = B.CreateGEP(EltTy, Ptr, Offset);
188 Value *Idx = B.CreateAdd(B.CreateMul(CurrentRow, B.getInt16(16)), CurrentCol);
189 if (IsTileLoad) {
190 // tileload.scalarize.rows.header:
191 // %vec.phi.row = phi <256 x i32> [ zeroinitializer, %entry ], [ %ResVec,
192 // %tileload.scalarize.rows.latch ]
193 B.SetInsertPoint(RowLoopHeader->getTerminator());
194 Value *VecZero = Constant::getNullValue(V256I32Ty);
195 PHINode *VecCPhiRowLoop = B.CreatePHI(V256I32Ty, 2, "vec.phi.row");
196 VecCPhiRowLoop->addIncoming(VecZero, Start);
197
198 // tileload.scalarize.cols.header:
199 // %vec.phi = phi <256 x i32> [ %vec.phi.row, %tileload.scalarize.rows.body
200 // ], [ %ResVec, %tileload.scalarize.cols.latch ]
201 B.SetInsertPoint(ColLoopHeader->getTerminator());
202 PHINode *VecPhi = B.CreatePHI(V256I32Ty, 2, "vec.phi");
203 VecPhi->addIncoming(VecCPhiRowLoop, RowBody);
204
205 // tileload.scalarize.cols.body:
206 // Calculate %idxmem and %idxvec
207 // %eltptr = getelementptr i32, i32* %base, i64 %idxmem
208 // %elt = load i32, i32* %ptr
209 // %ResVec = insertelement <256 x i32> %vec.phi, i32 %elt, i16 %idxvec
210 B.SetInsertPoint(ColBody->getTerminator());
211 Value *Elt = B.CreateLoad(EltTy, EltPtr);
212 Value *ResVec = B.CreateInsertElement(VecPhi, Elt, Idx);
213 VecPhi->addIncoming(ResVec, ColLoopLatch);
214 VecCPhiRowLoop->addIncoming(ResVec, RowLatch);
215
216 return ResVec;
217 } else {
218 auto *BitCast = cast<BitCastInst>(Tile);
219 Value *Vec = BitCast->getOperand(0);
220 assert(isV256I32Ty(Vec->getType()) && "bitcast from non-v256i32 to x86amx");
221 // tilestore.scalarize.cols.body:
222 // %mul = mul i16 %row.iv, i16 16
223 // %idx = add i16 %mul, i16 %col.iv
224 // %vec = extractelement <16 x i32> %vec, i16 %idx
225 // store i32 %vec, i32* %ptr
226 B.SetInsertPoint(ColBody->getTerminator());
227 Value *Elt = B.CreateExtractElement(Vec, Idx);
228
229 B.CreateStore(Elt, EltPtr);
230 return nullptr;
231 }
232}
233
234template <Intrinsic::ID IntrID>
235std::enable_if_t<IntrID == Intrinsic::x86_tdpbssd_internal ||
236 IntrID == Intrinsic::x86_tdpbsud_internal ||
237 IntrID == Intrinsic::x86_tdpbusd_internal ||
238 IntrID == Intrinsic::x86_tdpbuud_internal ||
239 IntrID == Intrinsic::x86_tdpbf16ps_internal,
240 Value *>
241X86LowerAMXIntrinsics::createTileDPLoops(BasicBlock *Start, BasicBlock *End,
242 IRBuilderBase &B, Value *Row,
243 Value *Col, Value *K, Value *Acc,
244 Value *LHS, Value *RHS) {
245 std::string IntrinName;
246 switch (IntrID) {
247 case Intrinsic::x86_tdpbssd_internal:
248 IntrinName = "tiledpbssd";
249 break;
250 case Intrinsic::x86_tdpbsud_internal:
251 IntrinName = "tiledpbsud";
252 break;
253 case Intrinsic::x86_tdpbusd_internal:
254 IntrinName = "tiledpbusd";
255 break;
256 case Intrinsic::x86_tdpbuud_internal:
257 IntrinName = "tiledpbuud";
258 break;
259 case Intrinsic::x86_tdpbf16ps_internal:
260 IntrinName = "tiledpbf16ps";
261 break;
262 }
263 Loop *RowLoop = nullptr;
264 Loop *ColLoop = nullptr;
265 Loop *InnerLoop = nullptr;
266 if (LI) {
267 RowLoop = LI->AllocateLoop();
268 ColLoop = LI->AllocateLoop();
269 InnerLoop = LI->AllocateLoop();
270 ColLoop->addChildLoop(InnerLoop);
271 RowLoop->addChildLoop(ColLoop);
272 if (Loop *ParentL = LI->getLoopFor(Start))
273 ParentL->addChildLoop(RowLoop);
274 else
275 LI->addTopLevelLoop(RowLoop);
276 }
277
278 BasicBlock *RowBody = createLoop(Start, End, Row, B.getInt16(1),
279 IntrinName + ".scalarize.rows", B, RowLoop);
280 BasicBlock *RowLatch = RowBody->getSingleSuccessor();
281
282 BasicBlock *ColBody = createLoop(RowBody, RowLatch, Col, B.getInt16(1),
283 IntrinName + ".scalarize.cols", B, ColLoop);
284
285 BasicBlock *ColLoopLatch = ColBody->getSingleSuccessor();
286
287 B.SetInsertPoint(ColBody->getTerminator());
288 BasicBlock *InnerBody =
289 createLoop(ColBody, ColLoopLatch, K, B.getInt16(1),
290 IntrinName + ".scalarize.inner", B, InnerLoop);
291
292 BasicBlock *ColLoopHeader = ColBody->getSinglePredecessor();
293 BasicBlock *RowLoopHeader = RowBody->getSinglePredecessor();
294 BasicBlock *InnerLoopHeader = InnerBody->getSinglePredecessor();
295 BasicBlock *InnerLoopLatch = InnerBody->getSingleSuccessor();
296 Value *CurrentRow = &*RowLoopHeader->begin();
297 Value *CurrentCol = &*ColLoopHeader->begin();
298 Value *CurrentInner = &*InnerLoopHeader->begin();
299
300 FixedVectorType *V256I32Ty = FixedVectorType::get(B.getInt32Ty(), 256);
301 auto *BitCastAcc = cast<BitCastInst>(Acc);
302 Value *VecC = BitCastAcc->getOperand(0);
303 assert(isV256I32Ty(VecC->getType()) && "bitcast from non-v256i32 to x86amx");
304 // TODO else create BitCast from x86amx to v256i32.
305 // Store x86amx to memory, and reload from memory
306 // to vector. However with -O0, it doesn't happen.
307 auto *BitCastLHS = cast<BitCastInst>(LHS);
308 Value *VecA = BitCastLHS->getOperand(0);
309 assert(isV256I32Ty(VecA->getType()) && "bitcast from non-v256i32 to x86amx");
310 auto *BitCastRHS = cast<BitCastInst>(RHS);
311 Value *VecB = BitCastRHS->getOperand(0);
312 assert(isV256I32Ty(VecB->getType()) && "bitcast from non-v256i32 to x86amx");
313
314 // tiledpbssd.scalarize.rows.header:
315 // %vec.c.phi.row = phi <256 x i32> [ %VecC, %continue ], [ %NewVecC,
316 // %tiledpbssd.scalarize.rows.latch ]
317
318 // %vec.d.phi.row = phi <256 x i32> [ zeroinitializer, %continue ], [
319 // %NewVecD, %tiledpbssd.scalarize.rows.latch ]
320 B.SetInsertPoint(RowLoopHeader->getTerminator());
321 PHINode *VecCPhiRowLoop = B.CreatePHI(V256I32Ty, 2, "vec.c.phi.row");
322 VecCPhiRowLoop->addIncoming(VecC, Start);
323 Value *VecZero = Constant::getNullValue(V256I32Ty);
324 PHINode *VecDPhiRowLoop = B.CreatePHI(V256I32Ty, 2, "vec.d.phi.row");
325 VecDPhiRowLoop->addIncoming(VecZero, Start);
326
327 // tiledpbssd.scalarize.cols.header:
328 // %vec.c.phi.col = phi <256 x i32> [ %vec.c.phi.row,
329 // %tiledpbssd.scalarize.rows.body ], [ %NewVecC,
330 // %tiledpbssd.scalarize.cols.latch ]
331
332 // %vec.d.phi.col = phi <256 x i32> [
333 // %vec.d.phi.row, %tiledpbssd.scalarize.rows.body ], [ %NewVecD,
334 // %tiledpbssd.scalarize.cols.latch ]
335
336 // calculate idxc.
337 B.SetInsertPoint(ColLoopHeader->getTerminator());
338 PHINode *VecCPhiColLoop = B.CreatePHI(V256I32Ty, 2, "vec.c.phi.col");
339 VecCPhiColLoop->addIncoming(VecCPhiRowLoop, RowBody);
340 PHINode *VecDPhiColLoop = B.CreatePHI(V256I32Ty, 2, "vec.d.phi.col");
341 VecDPhiColLoop->addIncoming(VecDPhiRowLoop, RowBody);
342 Value *IdxC =
343 B.CreateAdd(B.CreateMul(CurrentRow, B.getInt16(16)), CurrentCol);
344
345 // tiledpbssd.scalarize.inner.header:
346 // %vec.c.inner.phi = phi <256 x i32> [ %vec.c.phi.col,
347 // %tiledpbssd.scalarize.cols.body ], [ %NewVecC,
348 // %tiledpbssd.scalarize.inner.latch ]
349
350 B.SetInsertPoint(InnerLoopHeader->getTerminator());
351 PHINode *VecCPhi = B.CreatePHI(V256I32Ty, 2, "vec.c.inner.phi");
352 VecCPhi->addIncoming(VecCPhiColLoop, ColBody);
353
354 B.SetInsertPoint(InnerBody->getTerminator());
355 Value *IdxA =
356 B.CreateAdd(B.CreateMul(CurrentRow, B.getInt16(16)), CurrentInner);
357 Value *IdxB =
358 B.CreateAdd(B.CreateMul(CurrentInner, B.getInt16(16)), CurrentCol);
359 Value *NewVecC = nullptr;
360
361 if (IntrID != Intrinsic::x86_tdpbf16ps_internal) {
362 // tiledpbssd.scalarize.inner.body:
363 // calculate idxa, idxb
364 // %eltc = extractelement <256 x i32> %vec.c.inner.phi, i16 %idxc
365 // %elta = extractelement <256 x i32> %veca, i16 %idxa
366 // %eltav4i8 = bitcast i32 %elta to <4 x i8>
367 // %eltb = extractelement <256 x i32> %vecb, i16 %idxb
368 // %eltbv4i8 = bitcast i32 %eltb to <4 x i8>
369 // %eltav4i32 = sext <4 x i8> %eltav4i8 to <4 x i32>
370 // %eltbv4i32 = sext <4 x i8> %eltbv4i8 to <4 x i32>
371 // %mulab = mul <4 x i32> %eltbv4i32, %eltav4i32
372 // %acc = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> %131)
373 // %neweltc = add i32 %elt, %acc
374 // %NewVecC = insertelement <256 x i32> %vec.c.inner.phi, i32 %neweltc,
375 // i16 %idxc
376 FixedVectorType *V4I8Ty = FixedVectorType::get(B.getInt8Ty(), 4);
377 FixedVectorType *V4I32Ty = FixedVectorType::get(B.getInt32Ty(), 4);
378 Value *EltC = B.CreateExtractElement(VecCPhi, IdxC);
379 Value *EltA = B.CreateExtractElement(VecA, IdxA);
380 Value *SubVecA = B.CreateBitCast(EltA, V4I8Ty);
381 Value *EltB = B.CreateExtractElement(VecB, IdxB);
382 Value *SubVecB = B.CreateBitCast(EltB, V4I8Ty);
383 Value *SEXTSubVecB = nullptr;
384 Value *SEXTSubVecA = nullptr;
385 switch (IntrID) {
386 case Intrinsic::x86_tdpbssd_internal:
387 SEXTSubVecB = B.CreateSExt(SubVecB, V4I32Ty);
388 SEXTSubVecA = B.CreateSExt(SubVecA, V4I32Ty);
389 break;
390 case Intrinsic::x86_tdpbsud_internal:
391 SEXTSubVecB = B.CreateZExt(SubVecB, V4I32Ty);
392 SEXTSubVecA = B.CreateSExt(SubVecA, V4I32Ty);
393 break;
394 case Intrinsic::x86_tdpbusd_internal:
395 SEXTSubVecB = B.CreateSExt(SubVecB, V4I32Ty);
396 SEXTSubVecA = B.CreateZExt(SubVecA, V4I32Ty);
397 break;
398 case Intrinsic::x86_tdpbuud_internal:
399 SEXTSubVecB = B.CreateZExt(SubVecB, V4I32Ty);
400 SEXTSubVecA = B.CreateZExt(SubVecA, V4I32Ty);
401 break;
402 default:
403 llvm_unreachable("Invalid intrinsic ID!");
404 }
405 Value *SubVecR = B.CreateAddReduce(B.CreateMul(SEXTSubVecA, SEXTSubVecB));
406 Value *ResElt = B.CreateAdd(EltC, SubVecR);
407 NewVecC = B.CreateInsertElement(VecCPhi, ResElt, IdxC);
408 } else {
409 // tiledpbf16ps.scalarize.inner.body:
410 // calculate idxa, idxb, idxc
411 // %eltc = extractelement <256 x i32> %vec.c.inner.phi, i16 %idxc
412 // %eltcf32 = bitcast i32 %eltc to float
413 // %elta = extractelement <256 x i32> %veca, i16 %idxa
414 // %eltav2i16 = bitcast i32 %elta to <2 x i16>
415 // %eltb = extractelement <256 x i32> %vecb, i16 %idxb
416 // %eltbv2i16 = bitcast i32 %eltb to <2 x i16>
417 // %shufflea = shufflevector <2 x i16> %elta, <2 x i16> zeroinitializer, <4
418 // x i32> <i32 2, i32 0, i32 3, i32 1>
419 // %eltav2f32 = bitcast <4 x i16> %shufflea to <2 x float>
420 // %shuffleb = shufflevector <2 x i16> %eltb, <2 xi16> zeroinitializer, <4 x
421 // i32> <i32 2, i32 0, i32 3, i32 1>
422 // %eltbv2f32 = bitcast <4 x i16> %shuffleb to <2 x float>
423 // %mulab = fmul <2 x float> %eltav2f32, %eltbv2f32
424 // %acc = call float
425 // @llvm.vector.reduce.fadd.v2f32(float %eltcf32, <2 x float> %mulab)
426 // %neweltc = bitcast float %acc to i32
427 // %NewVecC = insertelement <256 x i32> %vec.c.inner.phi, i32 %neweltc,
428 // i16 %idxc
429 // %NewVecD = insertelement <256 x i32> %vec.d.inner.phi, i32 %neweltc,
430 // i16 %idxc
431 FixedVectorType *V2I16Ty = FixedVectorType::get(B.getInt16Ty(), 2);
432 FixedVectorType *V2F32Ty = FixedVectorType::get(B.getFloatTy(), 2);
433 Value *EltC = B.CreateExtractElement(VecCPhi, IdxC);
434 Value *EltCF32 = B.CreateBitCast(EltC, B.getFloatTy());
435 Value *EltA = B.CreateExtractElement(VecA, IdxA);
436 Value *SubVecA = B.CreateBitCast(EltA, V2I16Ty);
437 Value *EltB = B.CreateExtractElement(VecB, IdxB);
438 Value *SubVecB = B.CreateBitCast(EltB, V2I16Ty);
439 Value *ZeroV2I16 = Constant::getNullValue(V2I16Ty);
440 int ShuffleMask[4] = {2, 0, 3, 1};
441 auto ShuffleArray = ArrayRef(ShuffleMask);
442 Value *AV2F32 = B.CreateBitCast(
443 B.CreateShuffleVector(SubVecA, ZeroV2I16, ShuffleArray), V2F32Ty);
444 Value *BV2F32 = B.CreateBitCast(
445 B.CreateShuffleVector(SubVecB, ZeroV2I16, ShuffleArray), V2F32Ty);
446 Value *SubVecR = B.CreateFAddReduce(EltCF32, B.CreateFMul(AV2F32, BV2F32));
447 Value *ResElt = B.CreateBitCast(SubVecR, B.getInt32Ty());
448 NewVecC = B.CreateInsertElement(VecCPhi, ResElt, IdxC);
449 }
450
451 // tiledpbssd.scalarize.cols.latch:
452 // %NewEltC = extractelement <256 x i32> %vec.c.phi.col, i16 %idxc
453 // %NewVecD = insertelement <256 x i32> %vec.d.phi.col, i32 %NewEltC,
454 // i16 %idxc
455 B.SetInsertPoint(ColLoopLatch->getTerminator());
456 Value *NewEltC = B.CreateExtractElement(NewVecC, IdxC);
457 Value *NewVecD = B.CreateInsertElement(VecDPhiColLoop, NewEltC, IdxC);
458
459 VecCPhi->addIncoming(NewVecC, InnerLoopLatch);
460 VecCPhiRowLoop->addIncoming(NewVecC, RowLatch);
461 VecCPhiColLoop->addIncoming(NewVecC, ColLoopLatch);
462 VecDPhiRowLoop->addIncoming(NewVecD, RowLatch);
463 VecDPhiColLoop->addIncoming(NewVecD, ColLoopLatch);
464
465 return NewVecD;
466}
467
468template <Intrinsic::ID IntrID>
469std::enable_if_t<IntrID == Intrinsic::x86_tdpbssd_internal ||
470 IntrID == Intrinsic::x86_tdpbsud_internal ||
471 IntrID == Intrinsic::x86_tdpbusd_internal ||
472 IntrID == Intrinsic::x86_tdpbuud_internal ||
473 IntrID == Intrinsic::x86_tdpbf16ps_internal,
474 bool>
475X86LowerAMXIntrinsics::lowerTileDP(Instruction *TileDP) {
476 Value *M, *N, *K, *C, *A, *B;
477 match(TileDP, m_Intrinsic<IntrID>(m_Value(M), m_Value(N), m_Value(K),
478 m_Value(C), m_Value(A), m_Value(B)));
479 Instruction *InsertI = TileDP;
480 IRBuilder<> PreBuilder(TileDP);
481 PreBuilder.SetInsertPoint(TileDP);
482 // We visit the loop with (m, n/4, k/4):
483 // %n_dword = lshr i16 %n, 2
484 // %k_dword = lshr i16 %k, 2
485 Value *NDWord = PreBuilder.CreateLShr(N, PreBuilder.getInt16(2));
486 Value *KDWord = PreBuilder.CreateLShr(K, PreBuilder.getInt16(2));
487 BasicBlock *Start = InsertI->getParent();
488 BasicBlock *End =
489 SplitBlock(InsertI->getParent(), InsertI, &DTU, LI, nullptr, "continue");
490 IRBuilder<> Builder(TileDP);
491 Value *ResVec = createTileDPLoops<IntrID>(Start, End, Builder, M, NDWord,
492 KDWord, C, A, B);
493 // we cannot assume there always be bitcast after tiledpbssd. So we need to
494 // insert one bitcast as required
495 Builder.SetInsertPoint(End, End->getFirstNonPHIIt());
496 Value *ResAMX =
497 Builder.CreateBitCast(ResVec, Type::getX86_AMXTy(Builder.getContext()));
498 // Delete TileDP intrinsic and do some clean-up.
499 for (Use &U : llvm::make_early_inc_range(TileDP->uses())) {
500 Instruction *I = cast<Instruction>(U.getUser());
501 Value *Vec;
502 if (match(I, m_BitCast(m_Value(Vec)))) {
503 I->replaceAllUsesWith(ResVec);
504 I->eraseFromParent();
505 }
506 }
507 TileDP->replaceAllUsesWith(ResAMX);
508 TileDP->eraseFromParent();
509 return true;
510}
511
512template <bool IsTileLoad>
513bool X86LowerAMXIntrinsics::lowerTileLoadStore(Instruction *TileLoadStore) {
514 Value *M, *N, *Ptr, *Stride, *Tile;
515 if (IsTileLoad)
516 match(TileLoadStore,
517 m_Intrinsic<Intrinsic::x86_tileloadd64_internal>(
518 m_Value(M), m_Value(N), m_Value(Ptr), m_Value(Stride)));
519 else
520 match(TileLoadStore, m_Intrinsic<Intrinsic::x86_tilestored64_internal>(
521 m_Value(M), m_Value(N), m_Value(Ptr),
522 m_Value(Stride), m_Value(Tile)));
523
524 Instruction *InsertI = TileLoadStore;
525 IRBuilder<> PreBuilder(TileLoadStore);
526 PreBuilder.SetInsertPoint(TileLoadStore);
527 Value *NDWord = PreBuilder.CreateLShr(N, PreBuilder.getInt16(2));
528 Value *StrideDWord = PreBuilder.CreateLShr(Stride, PreBuilder.getInt64(2));
529 BasicBlock *Start = InsertI->getParent();
530 BasicBlock *End =
531 SplitBlock(InsertI->getParent(), InsertI, &DTU, LI, nullptr, "continue");
532 IRBuilder<> Builder(TileLoadStore);
533 Value *ResVec = createTileLoadStoreLoops<IsTileLoad>(
534 Start, End, Builder, M, NDWord, Ptr, StrideDWord,
535 IsTileLoad ? nullptr : Tile);
536 if (IsTileLoad) {
537 // we cannot assume there always be bitcast after tileload. So we need to
538 // insert one bitcast as required
539 Builder.SetInsertPoint(End, End->getFirstNonPHIIt());
540 Value *ResAMX =
541 Builder.CreateBitCast(ResVec, Type::getX86_AMXTy(Builder.getContext()));
542 // Delete tileloadd6 intrinsic and do some clean-up
543 for (Use &U : llvm::make_early_inc_range(TileLoadStore->uses())) {
544 Instruction *I = cast<Instruction>(U.getUser());
545 Value *Vec;
546 if (match(I, m_BitCast(m_Value(Vec)))) {
547 I->replaceAllUsesWith(ResVec);
548 I->eraseFromParent();
549 }
550 }
551 TileLoadStore->replaceAllUsesWith(ResAMX);
552 }
553 TileLoadStore->eraseFromParent();
554 return true;
555}
556
557bool X86LowerAMXIntrinsics::lowerTileZero(Instruction *TileZero) {
558 IRBuilder<> Builder(TileZero);
559 FixedVectorType *V256I32Ty = FixedVectorType::get(Builder.getInt32Ty(), 256);
560 Value *VecZero = Constant::getNullValue(V256I32Ty);
561 for (Use &U : llvm::make_early_inc_range(TileZero->uses())) {
562 Instruction *I = cast<Instruction>(U.getUser());
563 Value *Vec;
564 if (match(I, m_BitCast(m_Value(Vec)))) {
565 I->replaceAllUsesWith(VecZero);
566 I->eraseFromParent();
567 }
568 }
569 TileZero->eraseFromParent();
570 return true;
571}
572
573bool X86LowerAMXIntrinsics::visit() {
574 bool C = false;
575 SmallVector<IntrinsicInst *, 8> WorkList;
576 for (BasicBlock *BB : depth_first(&Func)) {
577 for (BasicBlock::iterator II = BB->begin(), IE = BB->end(); II != IE;) {
578 if (auto *Inst = dyn_cast<IntrinsicInst>(&*II++)) {
579 switch (Inst->getIntrinsicID()) {
580 case Intrinsic::x86_tdpbssd_internal:
581 case Intrinsic::x86_tdpbsud_internal:
582 case Intrinsic::x86_tdpbusd_internal:
583 case Intrinsic::x86_tdpbuud_internal:
584 case Intrinsic::x86_tileloadd64_internal:
585 case Intrinsic::x86_tilestored64_internal:
586 case Intrinsic::x86_tilezero_internal:
587 case Intrinsic::x86_tdpbf16ps_internal:
588 WorkList.push_back(Inst);
589 break;
590 default:
591 break;
592 }
593 }
594 }
595 }
596
597 for (auto *Inst : WorkList) {
598 switch (Inst->getIntrinsicID()) {
599 case Intrinsic::x86_tdpbssd_internal:
600 C = lowerTileDP<Intrinsic::x86_tdpbssd_internal>(Inst) || C;
601 break;
602 case Intrinsic::x86_tdpbsud_internal:
603 C = lowerTileDP<Intrinsic::x86_tdpbsud_internal>(Inst) || C;
604 break;
605 case Intrinsic::x86_tdpbusd_internal:
606 C = lowerTileDP<Intrinsic::x86_tdpbusd_internal>(Inst) || C;
607 break;
608 case Intrinsic::x86_tdpbuud_internal:
609 C = lowerTileDP<Intrinsic::x86_tdpbuud_internal>(Inst) || C;
610 break;
611 case Intrinsic::x86_tdpbf16ps_internal:
612 C = lowerTileDP<Intrinsic::x86_tdpbf16ps_internal>(Inst) || C;
613 break;
614 case Intrinsic::x86_tileloadd64_internal:
615 C = lowerTileLoadStore<true>(Inst) || C;
616 break;
617 case Intrinsic::x86_tilestored64_internal:
618 C = lowerTileLoadStore<false>(Inst) || C;
619 break;
620 case Intrinsic::x86_tilezero_internal:
621 C = lowerTileZero(Inst) || C;
622 break;
623 default:
624 llvm_unreachable("invalid amx intrinsics!");
625 }
626 }
627
628 return C;
629}
630
631namespace {
632class X86LowerAMXIntrinsicsLegacyPass : public FunctionPass {
633public:
634 static char ID;
635
636 X86LowerAMXIntrinsicsLegacyPass() : FunctionPass(ID) {
637 initializeX86LowerAMXIntrinsicsLegacyPassPass(
638 *PassRegistry::getPassRegistry());
639 }
640
641 bool runOnFunction(Function &F) override {
642 if (!X86ScalarizeAMX)
643 return false;
644 TargetMachine *TM = &getAnalysis<TargetPassConfig>().getTM<TargetMachine>();
645 if (!F.hasFnAttribute(Attribute::OptimizeNone) &&
646 TM->getOptLevel() != CodeGenOptLevel::None)
647 return false;
648
649 auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>();
650 auto *DT = DTWP ? &DTWP->getDomTree() : nullptr;
651 auto *LIWP = getAnalysisIfAvailable<LoopInfoWrapperPass>();
652 auto *LI = LIWP ? &LIWP->getLoopInfo() : nullptr;
653 DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
654
655 X86LowerAMXIntrinsics LAT(F, DTU, LI);
656 return LAT.visit();
657 }
658 StringRef getPassName() const override { return "Lower AMX intrinsics"; }
659
660 void getAnalysisUsage(AnalysisUsage &AU) const override {
661 AU.addPreserved<DominatorTreeWrapperPass>();
662 AU.addPreserved<LoopInfoWrapperPass>();
663 AU.addRequired<TargetPassConfig>();
664 }
665};
666} // namespace
667
668static const char PassName[] = "Lower AMX intrinsics";
669char X86LowerAMXIntrinsicsLegacyPass::ID = 0;
670INITIALIZE_PASS_BEGIN(X86LowerAMXIntrinsicsLegacyPass, DEBUG_TYPE, PassName,
671 false, false)
672INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
673INITIALIZE_PASS_END(X86LowerAMXIntrinsicsLegacyPass, DEBUG_TYPE, PassName,
674 false, false)
675
676FunctionPass *llvm::createX86LowerAMXIntrinsicsPass() {
677 return new X86LowerAMXIntrinsicsLegacyPass();
678}
B
static GCRegistry::Add< OcamlGC > B("ocaml", "ocaml 3.10-compatible GC")
A
static GCRegistry::Add< ErlangGC > A("erlang", "erlang-compatible garbage collector")
Idx
Returns the sub type a function will return at a given Idx Should correspond to the result type of an ExtractValue instruction executed with just that one unsigned Idx
Definition: DeadArgumentElimination.cpp:354
Name
std::string Name
Definition: ELFObjHandler.cpp:77
End
bool End
Definition: ELF_riscv.cpp:480
F
#define F(x, y, z)
Definition: MD5.cpp:55
I
#define I(x, y, z)
Definition: MD5.cpp:58
TM
const char LLVMTargetMachineRef TM
Definition: PassBuilderBindings.cpp:47
INITIALIZE_PASS_DEPENDENCY
#define INITIALIZE_PASS_DEPENDENCY(depName)
Definition: PassSupport.h:55
INITIALIZE_PASS_END
#define INITIALIZE_PASS_END(passName, arg, name, cfg, analysis)
Definition: PassSupport.h:59
INITIALIZE_PASS_BEGIN
#define INITIALIZE_PASS_BEGIN(passName, arg, name, cfg, analysis)
Definition: PassSupport.h:52
PostOrderIterator.h
This file builds on the ADT/GraphTraits.h file to build a generic graph post order iterator.
Cond
const SmallVectorImpl< MachineOperand > & Cond
Definition: RISCVRedundantCopyElimination.cpp:75
assert
assert(ImpDefSCC.getReg()==AMDGPU::SCC &&ImpDefSCC.isDef())
TargetPassConfig.h
Target-Independent Code Generator Pass Configuration Options pass.
TargetTransformInfo.h
This pass exposes codegen information to IR-level passes.
X86ScalarizeAMX
static cl::opt< bool > X86ScalarizeAMX("enable-x86-scalar-amx", cl::init(false), cl::Hidden, cl::desc("X86: enable AMX scalarizition."))
isV256I32Ty
static bool isV256I32Ty(Type *Ty)
Definition: X86LowerAMXIntrinsics.cpp:48
DEBUG_TYPE
#define DEBUG_TYPE
Definition: X86LowerAMXIntrinsics.cpp:45
PassName
static const char PassName[]
Definition: X86LowerAMXIntrinsics.cpp:668
IV
static const uint32_t IV[8]
Definition: blake3_impl.h:78
llvm::AnalysisUsage
Represent the analysis usage information of a pass.
Definition: PassAnalysisSupport.h:47
llvm::AnalysisUsage::addRequired
AnalysisUsage & addRequired()
Definition: PassAnalysisSupport.h:75
llvm::AnalysisUsage::addPreserved
AnalysisUsage & addPreserved()
Add the specified Pass class to the set of analyses preserved by this pass.
Definition: PassAnalysisSupport.h:98
llvm::ArrayRef
ArrayRef - Represent a constant reference to an array (0 or more elements consecutively in memory),...
Definition: ArrayRef.h:41
llvm::BasicBlock
LLVM Basic Block Representation.
Definition: BasicBlock.h:60
llvm::BasicBlock::begin
iterator begin()
Instruction iterator methods.
Definition: BasicBlock.h:430
llvm::BasicBlock::Create
static BasicBlock * Create(LLVMContext &Context, const Twine &Name="", Function *Parent=nullptr, BasicBlock *InsertBefore=nullptr)
Creates a new BasicBlock.
Definition: BasicBlock.h:199
llvm::BasicBlock::getSinglePredecessor
const BasicBlock * getSinglePredecessor() const
Return the predecessor of this block if it has a single predecessor block.
Definition: BasicBlock.cpp:452
llvm::BasicBlock::getSingleSuccessor
const BasicBlock * getSingleSuccessor() const
Return the successor of this block if it has a single successor.
Definition: BasicBlock.cpp:482
llvm::BasicBlock::getParent
const Function * getParent() const
Return the enclosing method, or null if none.
Definition: BasicBlock.h:206
llvm::BasicBlock::iterator
InstListType::iterator iterator
Instruction iterators...
Definition: BasicBlock.h:165
llvm::BasicBlock::getContext
LLVMContext & getContext() const
Get the context in which this basic block lives.
Definition: BasicBlock.cpp:168
llvm::BasicBlock::getTerminator
const Instruction * getTerminator() const LLVM_READONLY
Returns the terminator instruction if the block is well formed or null if the block is not well forme...
Definition: BasicBlock.h:221
llvm::BranchInst
Conditional or Unconditional Branch instruction.
Definition: Instructions.h:3439
llvm::BranchInst::Create
static BranchInst * Create(BasicBlock *IfTrue, BasicBlock::iterator InsertBefore)
Definition: Instructions.h:3500
llvm::BranchInst::getSuccessor
BasicBlock * getSuccessor(unsigned i) const
Definition: Instructions.h:3547
llvm::BranchInst::setSuccessor
void setSuccessor(unsigned idx, BasicBlock *NewSucc)
Definition: Instructions.h:3552
llvm::Constant::getNullValue
static Constant * getNullValue(Type *Ty)
Constructor to create a '0' constant of arbitrary type.
Definition: Constants.cpp:370
llvm::DominatorTreeWrapperPass
Legacy analysis pass which computes a DominatorTree.
Definition: Dominators.h:317
llvm::FixedVectorType
Class to represent fixed width SIMD vectors.
Definition: DerivedTypes.h:539
llvm::FixedVectorType::get
static FixedVectorType * get(Type *ElementType, unsigned NumElts)
Definition: Type.cpp:692
llvm::FunctionPass
FunctionPass class - This class is used to implement most global optimizations.
Definition: Pass.h:311
llvm::FunctionPass::runOnFunction
virtual bool runOnFunction(Function &F)=0
runOnFunction - Virtual method overriden by subclasses to do the per-function processing of the pass.
llvm::IRBuilderBase
Common base class shared among various IRBuilders.
Definition: IRBuilder.h:94
llvm::IRBuilder
This provides a uniform API for creating instructions and inserting them into a basic block: either a...
Definition: IRBuilder.h:2666
llvm::Instruction::getParent
const BasicBlock * getParent() const
Definition: Instruction.h:152
llvm::Instruction::eraseFromParent
InstListType::iterator eraseFromParent()
This method unlinks 'this' from the containing basic block and deletes it.
Definition: Instruction.cpp:104
llvm::LLVMContext
This is an important class for using LLVM in a threaded context.
Definition: LLVMContext.h:67
llvm::LoopBase::addChildLoop
void addChildLoop(LoopT *NewChild)
Add the specified loop to be a child of this loop.
Definition: GenericLoopInfo.h:387
llvm::LoopInfoWrapperPass
The legacy pass manager's analysis pass to compute loop information.
Definition: LoopInfo.h:593
llvm::Loop
Represents a single loop in the control flow graph.
Definition: LoopInfo.h:44
llvm::PHINode::addIncoming
void addIncoming(Value *V, BasicBlock *BB)
Add an incoming value to the end of the PHI list.
Definition: Instructions.h:3134
llvm::PHINode::Create
static PHINode * Create(Type *Ty, unsigned NumReservedValues, const Twine &NameStr, BasicBlock::iterator InsertBefore)
Constructors - NumReservedValues is a hint for the number of incoming edges that this phi node will h...
Definition: Instructions.h:3024
llvm::PassRegistry::getPassRegistry
static PassRegistry * getPassRegistry()
getPassRegistry - Access the global registry object, which is automatically initialized at applicatio...
Definition: PassRegistry.cpp:24
llvm::Pass::getAnalysisUsage
virtual void getAnalysisUsage(AnalysisUsage &) const
getAnalysisUsage - This function should be overriden by passes that need analysis information to do t...
Definition: Pass.cpp:98
llvm::Pass::getPassName
virtual StringRef getPassName() const
getPassName - Return a nice clean name for a pass.
Definition: Pass.cpp:81
llvm::SmallVectorTemplateBase::push_back
void push_back(const T &Elt)
Definition: SmallVector.h:426
llvm::SmallVector
This is a 'vector' (really, a variable-sized array), optimized for the case when the array is small.
Definition: SmallVector.h:1209
llvm::StringRef
StringRef - Represent a constant reference to a string, i.e.
Definition: StringRef.h:50
llvm::TargetMachine
Primary interface to the complete machine description for the target machine.
Definition: TargetMachine.h:76
llvm::TargetPassConfig
Target-Independent Code Generator Pass Configuration Options.
Definition: TargetPassConfig.h:85
llvm::Type
The instances of the Type class are immutable: once they are created, they are never changed.
Definition: Type.h:45
llvm::Type::getX86_AMXTy
static Type * getX86_AMXTy(LLVMContext &C)
llvm::Type::getInt16Ty
static IntegerType * getInt16Ty(LLVMContext &C)
llvm::Use
A Use represents the edge between a Value definition and its users.
Definition: Use.h:43
llvm::Value
LLVM Value Representation.
Definition: Value.h:74
llvm::Value::getType
Type * getType() const
All values are typed, get the type of this value.
Definition: Value.h:255
llvm::Value::replaceAllUsesWith
void replaceAllUsesWith(Value *V)
Change all uses of this to point to a new Value.
Definition: Value.cpp:534
llvm::Value::uses
iterator_range< use_iterator > uses()
Definition: Value.h:376
llvm_unreachable
#define llvm_unreachable(msg)
Marks that the current location is not supposed to be reachable.
Definition: ErrorHandling.h:143
llvm::CallingConv::ID
unsigned ID
LLVM IR allows to use arbitrary numbers as calling convention identifiers.
Definition: CallingConv.h:24
llvm::CallingConv::C
@ C
The default llvm calling convention, compatible with C.
Definition: CallingConv.h:34
llvm::PatternMatch::match
bool match(Val *V, const Pattern &P)
Definition: PatternMatch.h:49
llvm::PatternMatch::m_BitCast
CastOperator_match< OpTy, Instruction::BitCast > m_BitCast(const OpTy &Op)
Matches BitCast.
Definition: PatternMatch.h:1839
llvm::PatternMatch::m_Value
class_match< Value > m_Value()
Match an arbitrary value and ignore it.
Definition: PatternMatch.h:92
llvm::cl::init
initializer< Ty > init(const Ty &Val)
Definition: CommandLine.h:450
llvm
This is an optimization pass for GlobalISel generic memory operations.
Definition: AddressRanges.h:18
llvm::Offset
@ Offset
Definition: DWP.cpp:456
llvm::createX86LowerAMXIntrinsicsPass
FunctionPass * createX86LowerAMXIntrinsicsPass()
The pass transforms amx intrinsics to scalar operation if the function has optnone attribute or it is...
Definition: X86LowerAMXIntrinsics.cpp:676
llvm::initializeX86LowerAMXIntrinsicsLegacyPassPass
void initializeX86LowerAMXIntrinsicsLegacyPassPass(PassRegistry &)
llvm::make_early_inc_range
iterator_range< early_inc_iterator_impl< detail::IterOfRange< RangeT > > > make_early_inc_range(RangeT &&Range)
Make a range that does early increment to allow mutation of the underlying range without disrupting i...
Definition: STLExtras.h:656
llvm::SplitBlock
BasicBlock * SplitBlock(BasicBlock *Old, BasicBlock::iterator SplitPt, DominatorTree *DT, LoopInfo *LI=nullptr, MemorySSAUpdater *MSSAU=nullptr, const Twine &BBName="", bool Before=false)
Split the specified block at the specified instruction.
Definition: BasicBlockUtils.cpp:1083
llvm::depth_first
iterator_range< df_iterator< T > > depth_first(const T &G)
Definition: DepthFirstIterator.h:229
N
#define N
Generated on Wed Apr 24 2024 03:28:45 for LLVM by doxygen 1.9.6