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

Last change on this file since 6173 was 6055, checked in by cameron, 16 months ago

Various small fixes

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