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

Last change on this file since 6043 was 6043, checked in by xwa163, 11 months ago

Init checkin for lz4_grep count-only pipeline with multiplexing

File size: 10.8 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 = this->fill_address(b, 32, 4, b->CreateMul(b->CreateTrunc(swizzleIndex, b->getInt32Ty()),
123                                                                        b->getInt32(8)));
124
125                addresses = b->CreateAdd(addresses, nullAddress);
126
127                Value *gather_result = b->CreateCall(
128                        gatherFunc,
129                        {
130                                UndefValue::get(b->getBitBlockType()),
131                                b->CreatePointerCast(streamSetBlockBasePtr, b->getInt8PtrTy()),
132                                addresses,
133                                Constant::getAllOnesValue(b->getBitBlockType()),
134                                b->getInt8(1)
135                        }
136                );
137                // Shift the swizzle to the right to clear off any used bits ...
138                Value* unreadBitsVec = b->CreateLShr(gather_result, b->simd_fill(64, swizzleOffset));
139                // ... then to the left to align the bits with the buffer and combine them.
140                Value* pendingBitsVec = b->CreateShl(unreadBitsVec, b->simd_fill(64, updatedBufferSize));
141
142                bufferVecArray[iStreamSetIndex / 4] = b->CreateOr(updatedBufferVecArray[iStreamSetIndex / 4], pendingBitsVec);
143                updatedBufferVecArray[iStreamSetIndex / 4]->addIncoming(bufferVecArray[iStreamSetIndex / 4], enqueueBits);
144            }
145
146            // Update the buffer size with the number of bits we have actually enqueued
147            Value * const maxBufferSize = b->CreateAdd(b->CreateSub(PDEP_WIDTH, swizzleOffset), updatedBufferSize);
148            bufferSize = b->CreateUMin(maxBufferSize, PDEP_WIDTH);
149            updatedBufferSize->addIncoming(bufferSize, enqueueBits);
150
151            // ... and increment the source offset by the number we actually inserted
152            Value * const inserted = b->CreateSub(bufferSize, updatedBufferSize);
153            sourceOffset = b->CreateAdd(updatedSourceOffset, inserted);
154            updatedSourceOffset->addIncoming(sourceOffset, enqueueBits);
155
156            // INVESTIGATE: we can branch at most once here. I'm not sure whether the potential
157            // branch misprediction is better or worse than always filling from two swizzles to
158            // ensure that we have enough bits to deposit.
159            BasicBlock * const depositBits = b->CreateBasicBlock();
160            b->CreateUnlikelyCondBr(b->CreateICmpULT(bufferSize, required), enqueueBits, depositBits);
161
162            b->SetInsertPoint(depositBits);
163
164            // Apply PDEP to each element of the combined swizzle using the current PDEP mask
165            Value * const mask = b->CreateExtractElement(selectors, i);
166
167            for (int iStreamSetIndex = 0; iStreamSetIndex < mNumberOfStream; iStreamSetIndex += 4) {
168                Value * source_field = bufferVecArray[iStreamSetIndex / 4];
169                for (int iStreamIndex = iStreamSetIndex; iStreamIndex < iStreamSetIndex + 4; iStreamIndex++) {
170                    Value * PDEP_field = b->CreateCall(pdep, {b->CreateExtractElement(source_field, iStreamIndex - iStreamSetIndex), mask});
171                    resultArray[iStreamIndex] = b->CreateInsertElement(resultArray[iStreamIndex], PDEP_field, i);
172                }
173                bufferVecArray[iStreamSetIndex / 4] = b->CreateLShr(bufferVecArray[iStreamSetIndex / 4], b->simd_fill(64, required));
174            }
175
176            bufferSize = b->CreateSub(bufferSize, required);
177        }
178
179        for (int iStreamIndex = 0; iStreamIndex < mNumberOfStream; iStreamIndex++) {
180            // Store the result
181            Value * const outputStreamPtr = b->getOutputStreamBlockPtr("output", b->getSize(iStreamIndex), strideIndex);
182            b->CreateBlockAlignedStore(resultArray[iStreamIndex], outputStreamPtr);
183        }
184
185        BasicBlock * const finishedBlock = b->GetInsertBlock();
186        sourceOffsetPhi->addIncoming(sourceOffset, finishedBlock);
187        bufferSizePhi->addIncoming(bufferSize, finishedBlock);
188
189        for (int iStreamSetIndex = 0; iStreamSetIndex < mNumberOfStream; iStreamSetIndex += 4) {
190            bufferVecPhiArray[iStreamSetIndex / 4]->addIncoming(bufferVecArray[iStreamSetIndex / 4], finishedBlock);
191        }
192
193        Value * const nextStrideIndex = b->CreateAdd(strideIndex, b->getSize(1));
194        strideIndex->addIncoming(nextStrideIndex, finishedBlock);
195        b->CreateLikelyCondBr(b->CreateICmpNE(nextStrideIndex, numOfBlocks), processBlock, finishedStrides);
196
197        b->SetInsertPoint(finishedStrides);
198    }
199
200    llvm::Value* BitStreamGatherPDEPKernel::fill_address(const std::unique_ptr<kernel::KernelBuilder> & b, unsigned fw, unsigned field_count, Value * a) {
201        Type * singleFieldVecTy = VectorType::get(b->getIntNTy(fw), 1);
202        Value * aVec = b->CreateBitCast(a, singleFieldVecTy);
203        return b->CreateShuffleVector(aVec, UndefValue::get(singleFieldVecTy), Constant::getNullValue(VectorType::get(b->getInt32Ty(), field_count)));
204    }
205}
Note: See TracBrowser for help on using the repository browser.