source: icGREP/icgrep-devel/icgrep/kernels/bitstream_gather_pdep_kernel.cpp @ 6040

Last change on this file since 6040 was 6040, checked in by xwa163, 12 months ago

Init checkin for bitstream_pdep_kernel with gather intrinsics

File size: 10.6 KB
Line 
1
2#include "bitstream_gather_pdep_kernel.h"
3#include <kernels/kernel_builder.h>
4#include <llvm/Support/raw_ostream.h>
5#include <toolchain/toolchain.h>
6
7using namespace llvm;
8
9namespace kernel {
10
11    BitStreamGatherPDEPKernel::BitStreamGatherPDEPKernel(const std::unique_ptr<kernel::KernelBuilder> & b, const unsigned numberOfStream, std::string name)
12            : MultiBlockKernel(std::move(name),
13// input stream sets
14                               {Binding{b->getStreamSetTy(), "marker", FixedRate(), Principal()},
15                                Binding{b->getStreamSetTy(numberOfStream), "source", PopcountOf("marker")}},
16// output stream set
17                               {Binding{b->getStreamSetTy(numberOfStream), "output", FixedRate()}},
18                               {}, {}, {}),
19              mNumberOfStream(numberOfStream)
20             {
21
22    }
23
24    void BitStreamGatherPDEPKernel::generateMultiBlockLogic(const std::unique_ptr<KernelBuilder> & b, Value * const numOfBlocks) {
25        BasicBlock * const entry = b->GetInsertBlock();
26        BasicBlock * const processBlock = b->CreateBasicBlock("processBlock");
27        BasicBlock * const finishedStrides = b->CreateBasicBlock("finishedStrides");
28        const auto pdepWidth = 64; // TODO handle 32 bit machine
29        ConstantInt * const BLOCK_WIDTH = b->getSize(b->getBitBlockWidth());
30        ConstantInt * const PDEP_WIDTH = b->getSize(pdepWidth);
31
32        Function * pdep = nullptr;
33        if (pdepWidth == 64) {
34            pdep = Intrinsic::getDeclaration(b->getModule(), Intrinsic::x86_bmi_pdep_64);
35        } else if (pdepWidth == 32) {
36            pdep = Intrinsic::getDeclaration(b->getModule(), Intrinsic::x86_bmi_pdep_32);
37        } else {
38            report_fatal_error(getName() + ": PDEP width must be 32 or 64");
39        }
40
41        Constant * const ZERO = b->getSize(0);
42        Value * const sourceItemCount = b->getProcessedItemCount("source");
43
44        Value * const initialSourceOffset = b->CreateURem(sourceItemCount, BLOCK_WIDTH);
45        b->CreateBr(processBlock);
46
47        b->SetInsertPoint(processBlock);
48        PHINode * const strideIndex = b->CreatePHI(b->getSizeTy(), 2);
49        strideIndex->addIncoming(ZERO, entry);
50
51        std::vector<PHINode*> bufferVecPhiArray(mNumberOfStream / 4, NULL);
52        std::vector<Value*> bufferVecArray(mNumberOfStream / 4, NULL);
53        for (int iStreamSetIndex = 0; iStreamSetIndex < mNumberOfStream; iStreamSetIndex += 4) {
54            PHINode * const bufferPhi = b->CreatePHI(b->getBitBlockType(), 2, "bufferPhi");
55            bufferPhi->addIncoming(Constant::getNullValue(b->getBitBlockType()), entry);
56            bufferVecPhiArray[iStreamSetIndex / 4] = bufferPhi;
57            bufferVecArray[iStreamSetIndex / 4] = bufferPhi;
58        }
59
60
61        PHINode * const sourceOffsetPhi = b->CreatePHI(b->getSizeTy(), 2);
62        sourceOffsetPhi->addIncoming(initialSourceOffset, entry);
63        PHINode * const bufferSizePhi = b->CreatePHI(b->getSizeTy(), 2);
64        bufferSizePhi->addIncoming(ZERO, entry);
65
66        // Extract the values we will use in the main processing loop
67        Value * const markerStream = b->getInputStreamBlockPtr("marker", ZERO, strideIndex);
68        Value * const markerValue = b->CreateBlockAlignedLoad(markerStream);
69        Value * const selectors = b->fwCast(pdepWidth, markerValue);
70        Value * const numOfSelectors = b->simd_popcount(pdepWidth, selectors);
71
72        // For each element of the marker block
73        Value * bufferSize = bufferSizePhi;
74        Value * sourceOffset = sourceOffsetPhi;
75
76
77        std::vector<llvm::Value*> resultArray(mNumberOfStream, UndefValue::get(b->getBitBlockType()));
78        /**
79         * TODO Assumed that the bitblocktype is always <4 * i64>
80         *                    The i < 4 here comes from  ^
81         */
82        for (unsigned i = 0; i < 4; i++) {
83
84            // How many bits will we deposit?
85            Value * const required = b->CreateExtractElement(numOfSelectors, b->getSize(i));
86
87            // Aggressively enqueue any additional bits
88            BasicBlock * const entry = b->GetInsertBlock();
89            BasicBlock * const enqueueBits = b->CreateBasicBlock();
90            b->CreateBr(enqueueBits);
91
92            b->SetInsertPoint(enqueueBits);
93            PHINode * const updatedBufferSize = b->CreatePHI(bufferSize->getType(), 2);
94            updatedBufferSize->addIncoming(bufferSize, entry);
95            PHINode * const updatedSourceOffset = b->CreatePHI(sourceOffset->getType(), 2);
96            updatedSourceOffset->addIncoming(sourceOffset, entry);
97
98            std::vector<PHINode * > updatedBufferVecArray(mNumberOfStream / 4, NULL);
99            for (int iStreamSetIndex = 0; iStreamSetIndex < mNumberOfStream; iStreamSetIndex += 4) {
100                Value* buffer = bufferVecArray[iStreamSetIndex / 4];
101                PHINode * const updatedBuffer = b->CreatePHI(buffer->getType(), 2);
102                updatedBuffer->addIncoming(buffer, entry);
103                updatedBufferVecArray[iStreamSetIndex / 4] = updatedBuffer;
104            }
105
106            // Calculate the block and swizzle index of the current swizzle row
107            Value * const blockOffset = b->CreateUDiv(updatedSourceOffset, BLOCK_WIDTH);
108            Value * const swizzleIndex = b->CreateUDiv(b->CreateURem(updatedSourceOffset, BLOCK_WIDTH), PDEP_WIDTH);
109
110            Value * const swizzleOffset = b->CreateURem(updatedSourceOffset, PDEP_WIDTH);
111
112
113            for (int iStreamSetIndex = 0; iStreamSetIndex < mNumberOfStream; iStreamSetIndex += 4) {
114                // gather instruction can process 4 streams each time
115                Value *streamSetBlockBasePtr = b->getInputStreamBlockPtr("source", b->getSize(iStreamSetIndex),
116                                                                         blockOffset);
117
118                Function *gatherFunc = Intrinsic::getDeclaration(b->getModule(), Intrinsic::x86_avx2_gather_d_q_256);
119                Value *addresses = ConstantVector::get(
120                        {b->getInt32(0), b->getInt32(32), b->getInt32(64), b->getInt32(96)});
121
122                Value *nullAddress = ConstantVector::getNullValue(addresses->getType());
123                for (int i = 0; i < 4; i++) {
124                    nullAddress = b->CreateInsertElement(nullAddress,
125                                                         b->CreateMul(b->CreateTrunc(swizzleIndex, b->getInt32Ty()),
126                                                                      b->getInt32(8)), i);
127                };
128                addresses = b->CreateAdd(addresses, nullAddress);
129
130                Value *gather_result = b->CreateCall(
131                        gatherFunc,
132                        {
133                                UndefValue::get(b->getBitBlockType()),
134                                b->CreatePointerCast(streamSetBlockBasePtr, b->getInt8PtrTy()),
135                                addresses,
136                                Constant::getAllOnesValue(b->getBitBlockType()),
137                                b->getInt8(1)
138                        }
139                );
140                // Shift the swizzle to the right to clear off any used bits ...
141                Value* unreadBitsVec = b->CreateLShr(gather_result, b->simd_fill(64, swizzleOffset));
142                // ... then to the left to align the bits with the buffer and combine them.
143                Value* pendingBitsVec = b->CreateShl(unreadBitsVec, b->simd_fill(64, updatedBufferSize));
144
145                bufferVecArray[iStreamSetIndex / 4] = b->CreateOr(updatedBufferVecArray[iStreamSetIndex / 4], pendingBitsVec);
146                updatedBufferVecArray[iStreamSetIndex / 4]->addIncoming(bufferVecArray[iStreamSetIndex / 4], enqueueBits);
147            }
148
149            // Update the buffer size with the number of bits we have actually enqueued
150            Value * const maxBufferSize = b->CreateAdd(b->CreateSub(PDEP_WIDTH, swizzleOffset), updatedBufferSize);
151            bufferSize = b->CreateUMin(maxBufferSize, PDEP_WIDTH);
152            updatedBufferSize->addIncoming(bufferSize, enqueueBits);
153
154            // ... and increment the source offset by the number we actually inserted
155            Value * const inserted = b->CreateSub(bufferSize, updatedBufferSize);
156            sourceOffset = b->CreateAdd(updatedSourceOffset, inserted);
157            updatedSourceOffset->addIncoming(sourceOffset, enqueueBits);
158
159            // INVESTIGATE: we can branch at most once here. I'm not sure whether the potential
160            // branch misprediction is better or worse than always filling from two swizzles to
161            // ensure that we have enough bits to deposit.
162            BasicBlock * const depositBits = b->CreateBasicBlock();
163            b->CreateUnlikelyCondBr(b->CreateICmpULT(bufferSize, required), enqueueBits, depositBits);
164
165            b->SetInsertPoint(depositBits);
166
167            // Apply PDEP to each element of the combined swizzle using the current PDEP mask
168            Value * const mask = b->CreateExtractElement(selectors, i);
169
170            for (int iStreamSetIndex = 0; iStreamSetIndex < mNumberOfStream; iStreamSetIndex += 4) {
171                Value * source_field = bufferVecArray[iStreamSetIndex / 4];
172                for (int iStreamIndex = iStreamSetIndex; iStreamIndex < iStreamSetIndex + 4; iStreamIndex++) {
173                    Value * PDEP_field = b->CreateCall(pdep, {b->CreateExtractElement(source_field, iStreamIndex - iStreamSetIndex), mask});
174                    resultArray[iStreamIndex] = b->CreateInsertElement(resultArray[iStreamIndex], PDEP_field, i);
175                }
176                bufferVecArray[iStreamSetIndex / 4] = b->CreateLShr(bufferVecArray[iStreamSetIndex / 4], b->simd_fill(64, required));
177            }
178
179            bufferSize = b->CreateSub(bufferSize, required);
180        }
181
182        for (int iStreamIndex = 0; iStreamIndex < mNumberOfStream; iStreamIndex++) {
183            // Store the result
184            Value * const outputStreamPtr = b->getOutputStreamBlockPtr("output", b->getSize(iStreamIndex), strideIndex);
185            b->CreateBlockAlignedStore(resultArray[iStreamIndex], outputStreamPtr);
186        }
187
188        BasicBlock * const finishedBlock = b->GetInsertBlock();
189        sourceOffsetPhi->addIncoming(sourceOffset, finishedBlock);
190        bufferSizePhi->addIncoming(bufferSize, finishedBlock);
191
192        for (int iStreamSetIndex = 0; iStreamSetIndex < mNumberOfStream; iStreamSetIndex += 4) {
193            bufferVecPhiArray[iStreamSetIndex / 4]->addIncoming(bufferVecArray[iStreamSetIndex / 4], finishedBlock);
194        }
195
196        Value * const nextStrideIndex = b->CreateAdd(strideIndex, b->getSize(1));
197        strideIndex->addIncoming(nextStrideIndex, finishedBlock);
198        b->CreateLikelyCondBr(b->CreateICmpNE(nextStrideIndex, numOfBlocks), processBlock, finishedStrides);
199
200        b->SetInsertPoint(finishedStrides);
201    }
202
203}
Note: See TracBrowser for help on using the repository browser.