Halide/test/correctness/bit_counting.cpp at master · monkey-who/Halide · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

#include "Halide.h"
#include <stdio.h>
#include <stdint.h>
#include <string>

using namespace Halide;

uint32_t local_popcount(uint32_t v) {
    uint32_t count = 0;
    while (v) {
        if (v & 1) ++count;
        v >>= 1;
    }
    return count;
}

uint32_t local_count_trailing_zeros(uint32_t v) {
    for (uint32_t b = 0; b < 32; ++b) {
        if (v & (1 << b))
            // found a set bit
            return b;
    }
    return 0;
}

uint32_t local_count_leading_zeros(uint32_t v) {
    for (uint32_t b = 0; b < 32; ++b) {
        if (v & (1 << (31 - b)))
            // found a set bit
            return b;
    }
    return 0;
}

std::string as_bits(uint32_t v) {
    std::string ret;
    for (int i = 31; i >= 0; --i)
        ret += (v & (1 << i)) ? '1' : '0';
    return ret;
}

int main() {
    Image<uint32_t> input(256);
    for (int i = 0; i < 256; i++) {
        if (i < 16)
            input(i) = i;
        else if (i < 32)
            input(i) = 0xfffffffful - i;
        else
            input(i) = rand();
    }

    Var x;

    Func popcount_test("popcount_test");
    popcount_test(x) = popcount(input(x));

    Image<uint32_t> popcount_result = popcount_test.realize(256);
    for (int i = 0; i < 256; ++i) {
        if (popcount_result(i) != local_popcount(input(i))) {
            std::string bits_string = as_bits(input(i));
            printf("Popcount of %u [0b%s] returned %d (should be %d)\n",
                   input(i), bits_string.c_str(), popcount_result(i),
                   local_popcount(input(i)));
            return -1;
        }
    }

    Func ctlz_test("ctlz_test");
    ctlz_test(x) = count_leading_zeros(input(x));

    Image<uint32_t> ctlz_result = ctlz_test.realize(256);
    for (int i = 0; i < 256; ++i) {
        if (input(i) == 0)
            // results are undefined for zero input
            continue;

        if (ctlz_result(i) != local_count_leading_zeros(input(i))) {
            std::string bits_string = as_bits(input(i));
            printf("Ctlz of %u [0b%s] returned %d (should be %d)\n",
                   input(i), bits_string.c_str(), ctlz_result(i),
                   local_count_leading_zeros(input(i)));
            return -1;
        }
    }

    Func cttz_test("cttz_test");
    cttz_test(x) = count_trailing_zeros(input(x));

    Image<uint32_t> cttz_result = cttz_test.realize(256);
    for (int i = 0; i < 256; ++i) {
        if (input(i) == 0)
            // results are undefined for zero input
            continue;

        if (cttz_result(i) != local_count_trailing_zeros(input(i))) {
            std::string bits_string = as_bits(input(i));
            printf("Cttz of %u [0b%s] returned %d (should be %d)\n",
                   input(i), bits_string.c_str(), cttz_result(i),
                   local_count_trailing_zeros(input(i)));
            return -1;
        }
    }

    printf("Success!\n");
    return 0;
}