Revision c413e325e5481686b53586398673984df73fa57d authored by Andrew Adams on 23 February 2021, 22:26:49 UTC, committed by Andrew Adams on 23 February 2021, 22:26:49 UTC
gpu_reuse_shared_memory.cpp
#include "Halide.h"
using namespace Halide;
int multi_thread_type_test(MemoryType memory_type) {
Func f1("f1"), f2("f2"), f3("f3"), f4("f4"), f5("f5"), f6("f6");
Var x, y, z;
f1(x, y, z) = cast<uint8_t>(1);
f2(x, y, z) = cast<uint32_t>(f1(x + 1, y, z) + f1(x, y + 1, z));
f3(x, y, z) = cast<uint16_t>(f2(x + 1, y, z) + f2(x, y + 1, z));
f4(x, y, z) = cast<uint16_t>(f3(x + 1, y, z) + f3(x, y + 1, z));
f5(x, y, z) = cast<uint32_t>(f4(x + 1, y, z) + f4(x, y + 1, z));
f6(x, y, z) = cast<uint8_t>(f5(x + 1, y, z) + f5(x, y + 1, z));
Var thread_x, thread_y;
f6.compute_root().gpu_tile(x, y, thread_x, thread_y, 1, 1);
f5.compute_at(f6, x).gpu_threads(x, y).store_in(memory_type);
f4.compute_at(f6, x).gpu_threads(x, y).store_in(memory_type);
f3.compute_at(f6, x).gpu_threads(x, y).store_in(memory_type);
f2.compute_at(f6, x).gpu_threads(x, y).store_in(memory_type);
f1.compute_at(f6, x).gpu_threads(x, y).store_in(memory_type);
const int size_x = 200;
const int size_y = 200;
const int size_z = 4;
Buffer<uint8_t> out = f6.realize({size_x, size_y, size_z});
uint8_t correct = 32;
for (int z = 0; z < size_z; z++) {
for (int y = 0; y < size_y; y++) {
for (int x = 0; x < size_x; x++) {
if (out(x, y, z) != correct) {
printf("out(%d, %d, %d) = %d instead of %d\n",
x, y, z, out(x, y, z), correct);
return -1;
}
}
}
}
printf("OK\n");
return 0;
}
int pyramid_test(MemoryType memory_type) {
const int levels = 10;
const int size_x = 100;
const int size_y = 100;
Var x, y, z, xo, xi, yo, yi, thread_x, thread_y;
std::vector<Func> funcs(levels);
funcs[0](x, y) = 1;
for (int i = 1; i < levels; ++i) {
funcs[i](x, y) = funcs[i - 1](2 * x, y);
}
funcs[levels - 1]
.compute_root()
.gpu_tile(x, y, thread_x, thread_y, 3, 4);
for (int i = levels - 2; i >= 0; --i) {
funcs[i]
.compute_at(funcs[levels - 1], x)
.split(x, xo, xi, 1 << (levels - i - 1))
.gpu_threads(xo, y)
.store_in(memory_type);
}
Buffer<int> out = funcs[levels - 1].realize({size_x, size_y});
int correct = 1;
for (int y = 0; y < size_y; y++) {
for (int x = 0; x < size_x; x++) {
if (out(x, y) != correct) {
printf("out(%d, %d) = %d instead of %d\n",
x, y, out(x, y), correct);
return -1;
}
}
}
printf("OK\n");
return 0;
}
int inverted_pyramid_test(MemoryType memory_type) {
const int levels = 6;
const int size_x = 8 * 16 * 4;
const int size_y = 8 * 16 * 4;
Var x, y, z, yo, yi, yii, xo, xi, xii, thread_x, thread_y;
std::vector<Func> funcs(levels);
funcs[0](x, y) = 1;
for (int i = 1; i < levels; ++i) {
funcs[i](x, y) = funcs[i - 1](x / 2, y);
}
funcs[levels - 1]
.compute_root()
.tile(x, y, xi, yi, 64, 64)
.gpu_blocks(x, y)
.tile(xi, yi, xii, yii, 16, 16)
.gpu_threads(xi, yi);
for (int i = levels - 2; i >= 0; --i) {
funcs[i]
.compute_at(funcs[levels - 1], x)
.tile(x, y, xi, yi, 4, 4)
.gpu_threads(xi, yi)
.store_in(memory_type);
}
funcs[levels - 1]
.bound(x, 0, size_x)
.bound(y, 0, size_y);
Buffer<int> out = funcs[levels - 1].realize({size_x, size_y});
int correct = 1;
for (int y = 0; y < size_y; y++) {
for (int x = 0; x < size_x; x++) {
if (out(x, y) != correct) {
printf("out(%d, %d) = %d instead of %d\n",
x, y, out(x, y), correct);
return -1;
}
}
}
printf("OK\n");
return 0;
}
int dynamic_shared_test(MemoryType memory_type) {
Func f1, f2, f3, f4;
Var x, xo, xi, thread_xo;
f1(x) = x;
f2(x) = f1(x) + f1(2 * x);
f3(x) = f2(x) + f2(2 * x);
f4(x) = f3(x) + f3(2 * x);
f4.split(x, xo, xi, 16).gpu_tile(xo, thread_xo, 16);
f3.compute_at(f4, xo).split(x, xo, xi, 16).gpu_threads(xi).store_in(memory_type);
f2.compute_at(f4, xo).split(x, xo, xi, 16).gpu_threads(xi).store_in(memory_type);
f1.compute_at(f4, xo).split(x, xo, xi, 16).gpu_threads(xi).store_in(memory_type);
// The amount of shared memory required varies with x
Buffer<int> out = f4.realize({500});
for (int x = 0; x < out.width(); x++) {
int correct = 27 * x;
if (out(x) != correct) {
printf("out(%d) = %d instead of %d\n",
x, out(x), correct);
return -1;
}
}
printf("OK\n");
return 0;
}
int main(int argc, char **argv) {
Target t = get_jit_target_from_environment();
if (!t.has_gpu_feature()) {
printf("[SKIP] No GPU target enabled.\n");
return 0;
}
for (auto memory_type : {MemoryType::GPUShared, MemoryType::Heap}) {
printf("Running multi thread type test\n");
if (multi_thread_type_test(memory_type) != 0) {
return -1;
}
printf("Running pyramid test\n");
if (pyramid_test(memory_type) != 0) {
return -1;
}
printf("Running inverted pyramid test\n");
if (inverted_pyramid_test(memory_type) != 0) {
return -1;
}
printf("Running dynamic shared test\n");
if (t.has_feature(Target::OpenGLCompute) && memory_type == MemoryType::GPUShared) {
printf("Skipping test because GL doesn't support dynamic sizes for shared memory\n");
} else {
if (dynamic_shared_test(memory_type) != 0) {
return -1;
}
}
}
printf("Success!\n");
return 0;
}
![swh spinner](/static/img/swh-spinner.gif)
Computing file changes ...