1 #include "DeviceArgument.h"
2 #include "CodeGen_GPU_Dev.h"
3
4 namespace Halide {
5 namespace Internal {
6
HostClosure(const Stmt & s,const std::string & loop_variable)7 HostClosure::HostClosure(const Stmt &s, const std::string &loop_variable) {
8 if (!loop_variable.empty()) {
9 ignore.push(loop_variable);
10 }
11 s.accept(this);
12 }
13
arguments()14 std::vector<DeviceArgument> HostClosure::arguments() {
15 std::vector<DeviceArgument> res;
16 for (const auto &v : vars) {
17 debug(2) << "var: " << v.first << "\n";
18 res.emplace_back(v.first, false, v.second, 0);
19 }
20 for (const auto &b : buffers) {
21 debug(2) << "buffer: " << b.first << " " << b.second.size;
22 if (b.second.read) debug(2) << " (read)";
23 if (b.second.write) debug(2) << " (write)";
24 debug(2) << "\n";
25
26 DeviceArgument arg(b.first, true, b.second.type, b.second.dimensions, b.second.size);
27 arg.read = b.second.read;
28 arg.write = b.second.write;
29 res.push_back(arg);
30 }
31 return res;
32 }
33
visit(const Call * op)34 void HostClosure::visit(const Call *op) {
35 if (op->is_intrinsic(Call::glsl_texture_load) ||
36 op->is_intrinsic(Call::image_load) ||
37 op->is_intrinsic(Call::glsl_texture_store) ||
38 op->is_intrinsic(Call::image_store)) {
39
40 // The argument to the call is either a StringImm or a broadcasted
41 // StringImm if this is part of a vectorized expression
42
43 const StringImm *string_imm = op->args[0].as<StringImm>();
44 if (!string_imm) {
45 internal_assert(op->args[0].as<Broadcast>());
46 string_imm = op->args[0].as<Broadcast>()->value.as<StringImm>();
47 }
48
49 internal_assert(string_imm);
50
51 std::string bufname = string_imm->value;
52 Buffer &ref = buffers[bufname];
53 ref.type = op->type;
54 // TODO: do we need to set ref.dimensions?
55
56 if (op->is_intrinsic(Call::glsl_texture_load) ||
57 op->is_intrinsic(Call::image_load)) {
58 ref.read = true;
59 } else if (op->is_intrinsic(Call::glsl_texture_store) ||
60 op->is_intrinsic(Call::image_store)) {
61 ref.write = true;
62 }
63
64 // The Func's name and the associated .buffer are mentioned in the
65 // argument lists, but don't treat them as free variables.
66 ScopedBinding<> p1(ignore, bufname);
67 ScopedBinding<> p2(ignore, bufname + ".buffer");
68 Internal::Closure::visit(op);
69 } else {
70 Internal::Closure::visit(op);
71 }
72 }
73
visit(const For * loop)74 void HostClosure::visit(const For *loop) {
75 if (CodeGen_GPU_Dev::is_gpu_var(loop->name)) {
76 // The size of the threads and blocks is not part of the closure
77 ScopedBinding<> p(ignore, loop->name);
78 loop->body.accept(this);
79 } else {
80 Internal::Closure::visit(loop);
81 }
82 }
83
84 } // namespace Internal
85 } // namespace Halide
86