source: icGREP/icgrep-devel/icgrep/wc.cpp @ 5957

Last change on this file since 5957 was 5946, checked in by cameron, 18 months ago

Fix for wc

File size: 11.2 KB
Line 
1/*
2 *  Copyright (c) 2018 International Characters.
3 *  This software is licensed to the public under the Open Software License 3.0.
4 *  icgrep is a trademark of International Characters.
5 */
6
7#include <iostream>
8#include <iomanip>
9#include <vector>
10#include <string>
11#include <toolchain/toolchain.h>
12#include <llvm/IR/Function.h>
13#include <llvm/IR/Module.h>
14#include <llvm/Support/CommandLine.h>
15#include <llvm/Support/raw_ostream.h>
16#include <cc/cc_compiler.h>
17#include <pablo/pablo_kernel.h>
18#include <kernels/kernel_builder.h>
19#include <IR_Gen/idisa_target.h>
20#include <kernels/streamset.h>
21#include <kernels/source_kernel.h>
22#include <kernels/s2p_kernel.h>
23#include <pablo/pablo_compiler.h>
24#include <pablo/pablo_toolchain.h>
25#include <toolchain/cpudriver.h>
26#include <sys/stat.h>
27#include <fcntl.h>
28#include <util/file_select.h>
29
30using namespace llvm;
31
32static cl::OptionCategory wcFlags("Command Flags", "wc options");
33
34static cl::list<std::string> inputFiles(cl::Positional, cl::desc("<input file ...>"), cl::OneOrMore, cl::cat(wcFlags));
35
36std::vector<std::string> allFiles;
37
38enum CountOptions {
39    LineOption, WordOption, CharOption, ByteOption
40};
41
42static cl::list<CountOptions> wcOptions(
43  cl::values(clEnumValN(LineOption, "l", "Report the number of lines in each input file."),
44             clEnumValN(WordOption, "w", "Report the number of words in each input file."),
45             clEnumValN(CharOption, "m", "Report the number of characters in each input file (override -c)."),
46             clEnumValN(ByteOption, "c", "Report the number of bytes in each input file (override -m).")
47             CL_ENUM_VAL_SENTINEL), cl::cat(wcFlags), cl::Grouping);
48                                                 
49static std::string wc_modes = "";
50
51static int defaultDisplayColumnWidth = 7;  // default field width
52
53
54
55bool CountLines = false;
56bool CountWords = false;
57bool CountChars = false;
58bool CountBytes = false;
59
60std::vector<uint64_t> lineCount;
61std::vector<uint64_t> wordCount;
62std::vector<uint64_t> charCount;
63std::vector<uint64_t> byteCount;
64
65uint64_t TotalLines = 0;
66uint64_t TotalWords = 0;
67uint64_t TotalChars = 0;
68uint64_t TotalBytes = 0;
69
70using namespace pablo;
71using namespace kernel;
72using namespace parabix;
73
74//  The callback routine that records counts in progress.
75//
76extern "C" {
77    void record_counts(uint64_t lines, uint64_t words, uint64_t chars, uint64_t bytes, uint64_t fileIdx) {
78        lineCount[fileIdx] = lines;
79        wordCount[fileIdx] = words;
80        charCount[fileIdx] = chars;
81        byteCount[fileIdx] = bytes;
82        TotalLines += lines;
83        TotalWords += words;
84        TotalChars += chars;
85        TotalBytes += bytes;
86    }
87}
88
89class WordCountKernel final: public pablo::PabloKernel {
90public:
91    WordCountKernel(const std::unique_ptr<kernel::KernelBuilder> & b, Binding && inputStreamSet);
92    bool isCachable() const override { return true; }
93    bool hasSignature() const override { return false; }
94protected:
95    void generatePabloMethod() override;
96};
97
98WordCountKernel::WordCountKernel (const std::unique_ptr<kernel::KernelBuilder> & b, Binding && inputStreamSet)
99: PabloKernel(b, "wc_" + wc_modes,
100    {inputStreamSet},
101    {},
102    {},
103    {Binding{b->getSizeTy(), "lineCount"}, Binding{b->getSizeTy(), "wordCount"}, Binding{b->getSizeTy(), "charCount"}}) {
104
105}
106
107void WordCountKernel::generatePabloMethod() {
108    PabloBuilder pb(getEntryScope());
109    std::unique_ptr<cc::CC_Compiler> ccc;
110    if (CountWords || CountChars) {
111        ccc = make_unique<cc::Parabix_CC_Compiler>(getEntryScope(), getInputStreamSet("u8bit"));
112    } else {
113        ccc = make_unique<cc::Direct_CC_Compiler>(getEntryScope(), pb.createExtract(getInput(0), pb.getInteger(0)));
114    }
115
116    //  output: 3 counters
117    Var * lc = getOutputScalarVar("lineCount");
118    Var * wc = getOutputScalarVar("wordCount");
119    Var * cc = getOutputScalarVar("charCount");
120
121    if (CountLines) {
122        PabloAST * LF = ccc->compileCC(re::makeByte(0x0A));
123        pb.createAssign(lc, pb.createCount(LF));
124    }
125    if (CountWords) {
126        PabloAST * WS = ccc->compileCC(re::makeCC(re::makeByte(0x09, 0x0D), re::makeByte(0x20)));
127        PabloAST * wordChar = pb.createNot(WS);
128        // WS_follow_or_start = 1 past WS or at start of file
129        PabloAST * WS_follow_or_start = pb.createNot(pb.createAdvance(wordChar, 1));
130        PabloAST * wordStart = pb.createInFile(pb.createAnd(wordChar, WS_follow_or_start));
131        pb.createAssign(wc, pb.createCount(wordStart));
132    }
133    if (CountChars) {
134        //
135        // FIXME: This correctly counts characters assuming valid UTF-8 input.  But what if input is
136        // not UTF-8, or is not valid?
137        //
138        PabloAST * u8Begin = ccc->compileCC(re::makeCC(re::makeByte(0, 0x7F), re::makeByte(0xC2, 0xF4)));
139        pb.createAssign(cc, pb.createCount(u8Begin));
140    }
141}
142
143typedef void (*WordCountFunctionType)(uint32_t fd, size_t fileIdx);
144
145void wcPipelineGen(ParabixDriver & pxDriver) {
146
147    auto & iBuilder = pxDriver.getBuilder();
148    Module * m = iBuilder->getModule();
149    const unsigned segmentSize = codegen::SegmentSize;
150    const unsigned bufferSegments = codegen::ThreadNum+1;
151
152   
153    Type * const int32Ty = iBuilder->getInt32Ty();
154    Type * const sizeTy = iBuilder->getSizeTy();
155    Type * const voidTy = iBuilder->getVoidTy();
156
157    FunctionType * const recordCountsType = FunctionType::get(voidTy, {sizeTy, sizeTy, sizeTy, sizeTy, sizeTy}, false);
158    Constant * const recordCounts = m->getOrInsertFunction("record_counts", recordCountsType);
159
160    FunctionType * const mainType = FunctionType::get(voidTy, {int32Ty, sizeTy}, false);
161    Function * const main = cast<Function>(m->getOrInsertFunction("Main", mainType));
162    main->setCallingConv(CallingConv::C);
163    Function::arg_iterator args = main->arg_begin();   
164    Value * const fileDecriptor = &*(args++);
165    fileDecriptor->setName("fileDecriptor");
166    Value * const fileIdx = &*(args++);
167    fileIdx->setName("fileIdx");
168
169    iBuilder->SetInsertPoint(BasicBlock::Create(m->getContext(), "entry", main,0));
170
171    StreamSetBuffer * const ByteStream = pxDriver.addBuffer<SourceBuffer>(iBuilder, iBuilder->getStreamSetTy(1, 8));
172
173
174    Kernel * mmapK = pxDriver.addKernelInstance<MMapSourceKernel>(iBuilder);
175    mmapK->setInitialArguments({fileDecriptor});
176    pxDriver.makeKernelCall(mmapK, {}, {ByteStream});
177   
178    Kernel * wck  = nullptr;
179    if (CountWords || CountChars) {
180        StreamSetBuffer * const BasisBits = pxDriver.addBuffer<CircularBuffer>(iBuilder, iBuilder->getStreamSetTy(8, 1), segmentSize * bufferSegments);
181        Kernel * s2pk = pxDriver.addKernelInstance<S2PKernel>(iBuilder);
182        pxDriver.makeKernelCall(s2pk, {ByteStream}, {BasisBits});
183       
184        wck = pxDriver.addKernelInstance<WordCountKernel>(iBuilder, Binding{iBuilder->getStreamSetTy(8, 1), "u8bit"});
185        pxDriver.makeKernelCall(wck, {BasisBits}, {});
186
187
188    } else {
189        wck = pxDriver.addKernelInstance<WordCountKernel>(iBuilder, Binding{iBuilder->getStreamSetTy(1, 8), "u8byte"});
190        pxDriver.makeKernelCall(wck, {ByteStream}, {});
191    }
192
193    pxDriver.generatePipelineIR();
194   
195    iBuilder->setKernel(mmapK);
196    Value * const fileSize = iBuilder->getAccumulator("fileSize");
197    iBuilder->setKernel(wck);
198    Value * const lineCount = iBuilder->getAccumulator("lineCount");
199    Value * const wordCount = iBuilder->getAccumulator("wordCount");
200    Value * const charCount = iBuilder->getAccumulator("charCount");
201
202    iBuilder->CreateCall(recordCounts, {lineCount, wordCount, charCount, fileSize, fileIdx});
203    pxDriver.deallocateBuffers();
204    iBuilder->CreateRetVoid();
205
206    pxDriver.finalizeObject();
207}
208
209
210
211void wc(WordCountFunctionType fn_ptr, const int64_t fileIdx) {
212    std::string fileName = allFiles[fileIdx];
213    struct stat sb;
214    const int fd = open(fileName.c_str(), O_RDONLY);
215    if (LLVM_UNLIKELY(fd == -1)) {
216        if (errno == EACCES) {
217            std::cerr << "wc: " << fileName << ": Permission denied.\n";
218        }
219        else if (errno == ENOENT) {
220            std::cerr << "wc: " << fileName << ": No such file.\n";
221        }
222        else {
223            std::cerr << "wc: " << fileName << ": Failed.\n";
224        }
225        return;
226    }
227    if (stat(fileName.c_str(), &sb) == 0 && S_ISDIR(sb.st_mode)) {
228        std::cerr << "wc: " << fileName << ": Is a directory.\n";
229        close(fd);
230        return;
231    }
232    fn_ptr(fd, fileIdx);
233    close(fd);
234}
235
236int main(int argc, char *argv[]) {
237    codegen::ParseCommandLineOptions(argc, argv, {&wcFlags, pablo_toolchain_flags(), codegen::codegen_flags()});
238    if (argv::RecursiveFlag || argv::DereferenceRecursiveFlag) {
239        argv::DirectoriesFlag = argv::Recurse;
240    }
241    allFiles = argv::getFullFileList(inputFiles);
242   
243    const auto fileCount = allFiles.size();
244    if (wcOptions.size() == 0) {
245        CountLines = true;
246        CountWords = true;
247        CountBytes = true;
248    } else {
249        CountLines = false;
250        CountWords = false;
251        CountBytes = false;
252        CountChars = false;
253        for (unsigned i = 0; i < wcOptions.size(); i++) {
254            switch (wcOptions[i]) {
255                case WordOption: CountWords = true; break;
256                case LineOption: CountLines = true; break;
257                case CharOption: CountChars = true; CountBytes = false; break;
258                case ByteOption: CountBytes = true; CountChars = false; break;
259            }
260        }
261    }
262    if (CountLines) wc_modes += "l";
263    if (CountWords) wc_modes += "w";
264    if (CountChars) wc_modes += "m";
265    if (CountBytes) wc_modes += "c";
266
267    ParabixDriver pxDriver("wc");
268    wcPipelineGen(pxDriver);
269    auto wordCountFunctionPtr = reinterpret_cast<WordCountFunctionType>(pxDriver.getMain());
270    lineCount.resize(fileCount);
271    wordCount.resize(fileCount);
272    charCount.resize(fileCount);
273    byteCount.resize(fileCount);
274
275    for (unsigned i = 0; i < fileCount; ++i) {
276        wc(wordCountFunctionPtr, i);
277    }
278   
279    size_t maxCount = 0;
280    if (CountLines) maxCount = TotalLines;
281    if (CountWords) maxCount = TotalWords;
282    if (CountChars) maxCount = TotalChars;
283    if (CountBytes) maxCount = TotalBytes;
284   
285    int displayColumnWidth = std::to_string(maxCount).size() + 1;
286    if (displayColumnWidth < defaultDisplayColumnWidth) displayColumnWidth = defaultDisplayColumnWidth;
287
288    for (unsigned i = 0; i < fileCount; ++i) {
289        std::cout << std::setw(displayColumnWidth-1);
290        if (CountLines) {
291            std::cout << lineCount[i] << std::setw(displayColumnWidth);
292        }
293        if (CountWords) {
294            std::cout << wordCount[i] << std::setw(displayColumnWidth);
295        }
296        if (CountChars) {
297            std::cout << charCount[i] << std::setw(displayColumnWidth);
298        }
299        if (CountBytes) {
300            std::cout << byteCount[i];
301        }
302        std::cout << " " << allFiles[i] << std::endl;
303    }
304    if (inputFiles.size() > 1) {
305        std::cout << std::setw(displayColumnWidth-1);
306        if (CountLines) {
307            std::cout << TotalLines << std::setw(displayColumnWidth);
308        }
309        if (CountWords) {
310            std::cout << TotalWords << std::setw(displayColumnWidth);
311        }
312        if (CountChars) {
313            std::cout << TotalChars << std::setw(displayColumnWidth);
314        }
315        if (CountBytes) {
316            std::cout << TotalBytes;
317        }
318        std::cout << " total" << std::endl;
319    }
320
321    return 0;
322}
Note: See TracBrowser for help on using the repository browser.