Revision 248e8da6cc845fb4c7765a3719ebc653ace60e4f authored by Steven Johnson on 20 April 2022, 17:41:49 UTC, committed by Steven Johnson on 20 April 2022, 17:41:49 UTC
UnpackBuffers.cpp
#include "UnpackBuffers.h"
#include "IROperator.h"
#include "IRVisitor.h"
#include <map>
namespace Halide {
namespace Internal {
using std::map;
using std::pair;
using std::set;
using std::string;
using std::vector;
namespace {
struct BufferInfo {
Expr handle;
int dimensions;
};
class FindBufferSymbols : public IRVisitor {
using IRVisitor::visit;
void visit_param(const string &ref_name, const Parameter ¶m) {
if (param.defined() && param.is_buffer()) {
const string &name = param.name();
buffers[name] =
BufferInfo{Variable::make(type_of<halide_buffer_t *>(), name + ".buffer", param),
param.dimensions()};
}
}
void visit_buffer(const string &ref_name, const Buffer<> &buffer) {
if (buffer.defined()) {
const string &name = buffer.name();
buffers[name] =
BufferInfo{Variable::make(type_of<halide_buffer_t *>(), name + ".buffer", buffer),
buffer.dimensions()};
}
}
void visit(const Variable *op) override {
visit_param(op->name, op->param);
visit_buffer(op->name, op->image);
symbols.insert(op->name);
}
void visit(const Load *op) override {
visit_param(op->name, op->param);
visit_buffer(op->name, op->image);
symbols.insert(op->name);
IRVisitor::visit(op);
}
void visit(const Store *op) override {
visit_param(op->name, op->param);
symbols.insert(op->name);
IRVisitor::visit(op);
}
public:
set<string> symbols;
map<string, BufferInfo> buffers;
};
} // namespace
Stmt unpack_buffers(Stmt s) {
FindBufferSymbols finder;
s.accept(&finder);
vector<pair<string, Expr>> lets;
for (auto &p : finder.buffers) {
const string &name = p.first;
const BufferInfo &info = p.second;
vector<Expr> args = {info.handle};
string host_var = name;
Expr host_val = Call::make(type_of<void *>(), Call::buffer_get_host, args, Call::Extern);
lets.emplace_back(host_var, host_val);
string dev_var = name + ".device";
Expr dev_val = Call::make(type_of<uint64_t>(), Call::buffer_get_device, args, Call::Extern);
lets.emplace_back(dev_var, dev_val);
string dev_interface_var = name + ".device_interface";
Expr dev_interface_val = Call::make(type_of<const halide_device_interface_t *>(),
Call::buffer_get_device_interface, args, Call::Extern);
lets.emplace_back(dev_interface_var, dev_interface_val);
string type_code_var = name + ".type";
Expr type_code_val = Call::make(UInt(32), Call::buffer_get_type, args, Call::Extern);
lets.emplace_back(type_code_var, type_code_val);
string host_dirty_var = name + ".host_dirty";
Expr host_dirty_val = Call::make(Bool(), Call::buffer_get_host_dirty, args, Call::Extern);
lets.emplace_back(host_dirty_var, host_dirty_val);
string dev_dirty_var = name + ".device_dirty";
Expr dev_dirty_val = Call::make(Bool(), Call::buffer_get_device_dirty, args, Call::Extern);
lets.emplace_back(dev_dirty_var, dev_dirty_val);
string dimensions_var = name + ".dimensions";
Expr dimensions_val = Call::make(Int(32), Call::buffer_get_dimensions, args, Call::Extern);
lets.emplace_back(dimensions_var, dimensions_val);
for (int i = 0; i < info.dimensions; i++) {
vector<Expr> args = {info.handle, i};
string min_var = name + ".min." + std::to_string(i);
Expr min_val = Call::make(Int(32), Call::buffer_get_min, args, Call::Extern);
lets.emplace_back(min_var, min_val);
string extent_var = name + ".extent." + std::to_string(i);
Expr extent_val = Call::make(Int(32), Call::buffer_get_extent, args, Call::Extern);
lets.emplace_back(extent_var, extent_val);
string stride_var = name + ".stride." + std::to_string(i);
Expr stride_val = Call::make(Int(32), Call::buffer_get_stride, args, Call::Extern);
lets.emplace_back(stride_var, stride_val);
}
}
while (!lets.empty()) {
pair<string, Expr> l = lets.back();
lets.pop_back();
if (finder.symbols.count(l.first)) {
s = LetStmt::make(l.first, l.second, s);
}
}
// Create buffer is not null assertions
for (auto &p : finder.buffers) {
Expr buf = p.second.handle;
Expr cond = reinterpret<uint64_t>(buf) != 0;
Expr error = Call::make(Int(32), "halide_error_buffer_argument_is_null",
{p.first}, Call::Extern);
Stmt check = AssertStmt::make(cond, error);
s = Block::make(check, s);
}
return s;
}
} // namespace Internal
} // namespace Halide
Computing file changes ...