https://github.com/halide/Halide
Tip revision: e26ce62fc9011e0c2caf2bee0985e4d71f44086f authored by Andrew Adams on 16 October 2023, 17:15:15 UTC
Merge remote-tracking branch 'origin/main' into abadams/fix_7892
Merge remote-tracking branch 'origin/main' into abadams/fix_7892
Tip revision: e26ce62
CodeGen_OpenGLCompute_Dev.cpp
#include "CodeGen_OpenGLCompute_Dev.h"
#include "CSE.h"
#include "CodeGen_C.h"
#include "CodeGen_GPU_Dev.h"
#include "Debug.h"
#include "Deinterleave.h"
#include "FindIntrinsics.h"
#include "IRMatch.h"
#include "IRMutator.h"
#include "IROperator.h"
#include "Simplify.h"
#include <iomanip>
#include <limits>
#include <map>
namespace Halide {
namespace Internal {
using std::ostringstream;
using std::string;
using std::vector;
namespace {
char get_lane_suffix(int i) {
internal_assert(i >= 0 && i < 4);
return "rgba"[i];
}
class CodeGen_OpenGLCompute_C : public CodeGen_C {
public:
CodeGen_OpenGLCompute_C(std::ostream &s, const Target &t);
void add_kernel(const Stmt &stmt,
const std::string &name,
const std::vector<DeviceArgument> &args);
protected:
Type map_type(const Type &);
std::string print_name(const std::string &name) override;
std::string print_type(Type type, AppendSpaceIfNeeded space_option = DoNotAppendSpace) override;
using CodeGen_C::visit;
void visit(const Cast *) override;
void visit(const FloatImm *) override;
void visit(const UIntImm *) override;
void visit(const IntImm *) override;
void visit(const Max *op) override;
void visit(const Min *op) override;
void visit(const Mod *) override;
// these have specific functions
// in GLSL that operate on vectors
void visit(const EQ *) override;
void visit(const NE *) override;
void visit(const LT *) override;
void visit(const LE *) override;
void visit(const GT *) override;
void visit(const GE *) override;
void visit(const Shuffle *) override;
void visit(const For *) override;
void visit(const Ramp *op) override;
void visit(const Broadcast *op) override;
void visit(const Load *op) override;
void visit(const Store *op) override;
void visit(const Call *op) override;
void visit(const Allocate *op) override;
void visit(const Free *op) override;
void visit(const Select *op) override;
void visit(const Evaluate *op) override;
const std::map<std::string, std::string> builtin = {
{"abs", "abs"},
{"abs_f32", "abs"},
{"acos_f32", "acos"},
{"acosh_f32", "acosh"},
{"asin_f32", "asin"},
{"asinh_f32", "asinh"},
{"atan2_f32", "atan"}, // also called atan in GLSL
{"atan_f32", "atan"},
{"atanh_f32", "atanh"},
{"ceil_f32", "ceil"},
{"cos_f32", "cos"},
{"cosh_f32", "cosh"},
{"equal", "equal"},
{"exp_f32", "exp"},
{"fast_inverse_sqrt_f32", "inversesqrt"},
{"floor_f32", "floor"},
{"greaterThan", "greaterThan"},
{"greaterThanEqual", "greaterThanEqual"},
{"isnan", "isnan"},
{"lessThan", "lessThan"},
{"lessThanEqual", "lessThanEqual"},
{"log_f32", "log"},
{"max", "max"},
{"min", "min"},
{"mix", "mix"},
{"mod", "mod"},
{"notEqual", "notEqual"},
{"sin_f32", "sin"},
{"sinh_f32", "sinh"},
{"sqrt_f32", "sqrt"},
{"tan_f32", "tan"},
{"tanh_f32", "tanh"},
{"trunc_f32", "trunc"},
};
int workgroup_size[3] = {0, 0, 0};
// Maps each buffer with whether its base type is a vector.
std::map<string, bool> buffer_is_vector;
};
CodeGen_OpenGLCompute_C::CodeGen_OpenGLCompute_C(std::ostream &s, const Target &t)
: CodeGen_C(s, t) {
}
// Maps Halide types to appropriate GLSL types or emit error if no equivalent
// type is available.
Type CodeGen_OpenGLCompute_C::map_type(const Type &type) {
Type result = type;
if (type.is_scalar()) {
if (type.is_float()) {
user_assert(type.bits() <= 32)
<< "GLSL: Can't represent a float with " << type.bits() << " bits.\n";
result = Float(32);
} else if (type.is_bool()) {
// unchanged
} else if (type.is_int() && type.bits() <= 32) {
result = Int(32);
} else if (type.is_uint() && type.bits() <= 32) {
result = UInt(32);
} else {
user_error << "GLSL: Can't represent type '" << type << "'.\n";
}
} else {
user_assert(type.lanes() <= 4)
<< "GLSL: vector types wider than 4 aren't supported\n";
user_assert(type.is_bool() || type.is_int() || type.is_uint() || type.is_float())
<< "GLSL: Can't represent vector type '" << type << "'.\n";
Type scalar_type = type.element_of();
result = map_type(scalar_type).with_lanes(type.lanes());
}
return result;
}
// Identifiers containing double underscores '__' are reserved in GLSL, so we
// have to use a different name mangling scheme than in the C code generator.
string CodeGen_OpenGLCompute_C::print_name(const string &name) {
const string mangled = CodeGen_C::print_name(name);
return replace_all(mangled, "__", "XX");
}
string CodeGen_OpenGLCompute_C::print_type(Type type, AppendSpaceIfNeeded space) {
ostringstream oss;
type = map_type(type);
if (type.is_scalar()) {
if (type.is_float()) {
oss << "float";
} else if (type.is_bool()) {
oss << "bool";
} else if (type.is_int()) {
oss << "int";
} else if (type.is_uint()) {
oss << "uint";
} else {
internal_error << "GLSL: invalid type '" << type << "' encountered.\n";
}
} else {
if (type.is_float()) {
// no prefix for float vectors
} else if (type.is_bool()) {
oss << "b";
} else if (type.is_int()) {
oss << "i";
} else if (type.is_uint()) {
oss << "u";
} else {
internal_error << "GLSL: invalid type '" << type << "' encountered.\n";
}
oss << "vec" << type.lanes();
}
if (space == AppendSpace) {
oss << " ";
}
return oss.str();
}
string simt_intrinsic(const string &name) {
if (ends_with(name, ".__thread_id_x")) {
return "gl_LocalInvocationID.x";
} else if (ends_with(name, ".__thread_id_y")) {
return "gl_LocalInvocationID.y";
} else if (ends_with(name, ".__thread_id_z")) {
return "gl_LocalInvocationID.z";
} else if (ends_with(name, ".__thread_id_w")) {
internal_error << "4-dimension loops with " << name << " are not supported\n";
} else if (ends_with(name, ".__block_id_x")) {
return "gl_WorkGroupID.x";
} else if (ends_with(name, ".__block_id_y")) {
return "gl_WorkGroupID.y";
} else if (ends_with(name, ".__block_id_z")) {
return "gl_WorkGroupID.z";
} else if (ends_with(name, ".__block_id_w")) {
internal_error << "4-dimension loops with " << name << " are not supported\n";
}
internal_error << "simt_intrinsic called on bad variable name: " << name << "\n";
return "";
}
int thread_loop_workgroup_index(const string &name) {
string ids[] = {".__thread_id_x",
".__thread_id_y",
".__thread_id_z",
".__thread_id_w"};
for (size_t i = 0; i < sizeof(ids) / sizeof(string); i++) {
if (ends_with(name, ids[i])) {
return i;
}
}
return -1;
}
void CodeGen_OpenGLCompute_C::visit(const FloatImm *op) {
ostringstream oss;
// Print integral numbers with trailing ".0". For fractional numbers use a
// precision of 9 digits, which should be enough to recover the binary
// float unambiguously from the decimal representation (if iostreams
// implements correct rounding).
const float truncated = (op->value < 0 ? std::ceil(op->value) : std::floor(op->value));
if (truncated == op->value) {
oss << std::fixed << std::setprecision(1) << op->value;
} else {
oss << std::setprecision(9) << op->value;
}
id = oss.str();
}
void CodeGen_OpenGLCompute_C::visit(const UIntImm *op) {
if (op->type == Bool()) {
if (op->value == 1) {
id = "true";
} else {
id = "false";
}
} else {
id = std::to_string(op->value) + "u";
}
}
void CodeGen_OpenGLCompute_C::visit(const Max *op) {
print_expr(Call::make(op->type, "max", {op->a, op->b}, Call::PureExtern));
}
void CodeGen_OpenGLCompute_C::visit(const Min *op) {
print_expr(Call::make(op->type, "min", {op->a, op->b}, Call::PureExtern));
}
void CodeGen_OpenGLCompute_C::visit(const Mod *op) {
if (op->type.is_int() || op->type.is_uint()) {
// Just exploit the Euclidean identity
// FIXME: Why doesn't lower_euclidean_mod work for glsl?
// https://github.com/halide/Halide/issues/4979
Expr zero = make_zero(op->type);
Expr equiv = select(op->a == zero, zero,
op->a - (op->a / op->b) * op->b);
equiv = common_subexpression_elimination(equiv);
print_expr(equiv);
} else {
print_expr(Call::make(op->type, "mod", {op->a, op->b}, Call::Extern));
}
}
// The following comparisons are defined for ivec and vec
// types, so we don't use call_builtin
void CodeGen_OpenGLCompute_C::visit(const EQ *op) {
if (op->type.is_vector()) {
print_expr(Call::make(op->type, "equal", {op->a, op->b}, Call::Extern));
} else {
CodeGen_C::visit(op);
}
}
void CodeGen_OpenGLCompute_C::visit(const NE *op) {
if (op->type.is_vector()) {
print_expr(Call::make(op->type, "notEqual", {op->a, op->b}, Call::Extern));
} else {
CodeGen_C::visit(op);
}
}
void CodeGen_OpenGLCompute_C::visit(const LT *op) {
if (op->type.is_vector()) {
print_expr(Call::make(op->type, "lessThan", {op->a, op->b}, Call::Extern));
} else {
CodeGen_C::visit(op);
}
}
void CodeGen_OpenGLCompute_C::visit(const LE *op) {
if (op->type.is_vector()) {
print_expr(Call::make(op->type, "lessThanEqual", {op->a, op->b}, Call::Extern));
} else {
CodeGen_C::visit(op);
}
}
void CodeGen_OpenGLCompute_C::visit(const GT *op) {
if (op->type.is_vector()) {
print_expr(Call::make(op->type, "greaterThan", {op->a, op->b}, Call::Extern));
} else {
CodeGen_C::visit(op);
}
}
void CodeGen_OpenGLCompute_C::visit(const GE *op) {
if (op->type.is_vector()) {
print_expr(Call::make(op->type, "greaterThanEqual", {op->a, op->b}, Call::Extern));
} else {
CodeGen_C::visit(op);
}
}
void CodeGen_OpenGLCompute_C::visit(const Shuffle *op) {
// The halide Shuffle represents the llvm intrinisc
// shufflevector, however, for GLSL its use is limited to swizzling
// up to a four channel vec type.
internal_assert(op->vectors.size() == 1);
int shuffle_lanes = op->type.lanes();
internal_assert(shuffle_lanes <= 4);
string expr = print_expr(op->vectors[0]);
// Create a swizzle expression for the shuffle
string swizzle;
for (int i = 0; i != shuffle_lanes; ++i) {
int channel = op->indices[i];
internal_assert(channel < 4) << "Shuffle of invalid channel";
swizzle += get_lane_suffix(channel);
}
print_assignment(op->type, expr + "." + swizzle);
}
void CodeGen_OpenGLCompute_C::visit(const Call *op) {
if (op->is_intrinsic(Call::gpu_thread_barrier)) {
internal_assert(op->args.size() == 1) << "gpu_thread_barrier() intrinsic must specify memory fence type.\n";
const auto *fence_type_ptr = as_const_int(op->args[0]);
internal_assert(fence_type_ptr) << "gpu_thread_barrier() parameter is not a constant integer.\n";
auto fence_type = *fence_type_ptr;
stream << get_indent() << "barrier();\n";
// barrier() is an execution barrier; for memory behavior, we'll use the
// least-common-denominator groupMemoryBarrier(), because other fence types
// require extensions or GL 4.3 as a minumum.
if (fence_type & CodeGen_GPU_Dev::MemoryFenceType::Device ||
fence_type & CodeGen_GPU_Dev::MemoryFenceType::Shared) {
stream << "groupMemoryBarrier();\n";
}
print_assignment(op->type, "0");
} else if (op->is_intrinsic(Call::lerp)) {
// Implement lerp using GLSL's mix() function, which always uses
// floating point arithmetic.
Expr zero_val = op->args[0];
Expr one_val = op->args[1];
Expr weight = op->args[2];
internal_assert(weight.type().is_uint() || weight.type().is_float());
if (weight.type().is_uint()) {
// Normalize integer weights to [0.0f, 1.0f] range.
internal_assert(weight.type().bits() < 32);
weight = Div::make(Cast::make(Float(32), weight),
Cast::make(Float(32), weight.type().max()));
} else if (op->type.is_uint()) {
// Round float weights down to next multiple of (1/op->type.imax())
// to give same results as lerp based on integer arithmetic.
internal_assert(op->type.bits() < 32);
weight = floor(weight * op->type.max()) / op->type.max();
}
Type result_type = Float(32, op->type.lanes());
Expr e = Call::make(result_type, "mix", {zero_val, one_val, weight}, Call::Extern);
if (!op->type.is_float()) {
// Mirror rounding implementation of Halide's integer lerp.
e = Cast::make(op->type, floor(e + 0.5f));
}
print_expr(e);
return;
} else if (op->is_intrinsic(Call::abs)) {
internal_assert(op->args.size() == 1);
Expr a = op->args[0];
Type target_type = map_type(op->type);
if (op->type != Int(32)) {
print_assignment(target_type, print_type(target_type) + "(abs(" + print_expr(a) + "))");
} else {
print_assignment(target_type, "abs(" + print_expr(a) + ")");
}
return;
} else if (op->is_intrinsic(Call::absd)) {
internal_assert(op->args.size() == 2);
Expr a = op->args[0];
Expr b = op->args[1];
Expr e = cast(op->type, select(a < b, b - a, a - b));
print_expr(e);
return;
} else if (op->is_intrinsic(Call::return_second)) {
internal_assert(op->args.size() == 2);
// Simply discard the first argument, which is generally a call to
// 'halide_printf'.
print_assignment(op->type, print_expr(op->args[1]));
return;
} else if (op->is_intrinsic(Call::round)) {
print_assignment(op->type, "roundEven(" + print_expr(op->args[0]) + ")");
return;
} else if (op->name == "fast_inverse_f32") {
print_expr(make_one(op->type) / op->args[0]);
return;
} else if (op->name == "fast_inverse_sqrt_f32") {
print_expr(make_one(op->type) / sqrt(op->args[0]));
return;
} else if (op->name == "pow_f32") {
if (can_prove(op->args[0] > 0)) {
ostringstream rhs;
rhs << "pow(" << print_expr(op->args[0]) << ", " << print_expr(op->args[1]) << ")";
print_assignment(op->type, rhs.str());
return;
} else {
ostringstream base;
string a = print_expr(op->args[0]);
string b = print_expr(op->args[1]);
base << "pow(abs(" << a << "), " << b << ")";
string c = print_assignment(op->type, base.str());
Expr a_var = is_const(op->args[0]) ? op->args[0] : Variable::make(op->type, a);
Expr b_var = is_const(op->args[1]) ? op->args[1] : Variable::make(op->type, b);
Expr c_var = Variable::make(op->type, c);
// OpenGL isn't required to produce NaNs, so we return
// zero in the undefined case.
Expr equiv = select(a_var > 0 || b_var % 2 == 0, c_var,
b_var % 2 == 1, -c_var,
0.0f);
print_expr(simplify(equiv));
return;
}
} else if (op->is_intrinsic(Call::shift_right)) {
print_assignment(op->type, print_expr(op->args[0]) + " >> " + print_expr(op->args[1]));
} else if (op->is_intrinsic(Call::shift_left)) {
print_assignment(op->type, print_expr(op->args[0]) + " << " + print_expr(op->args[1]));
} else if (op->is_intrinsic(Call::bitwise_not)) {
print_assignment(op->type, "~" + print_expr(op->args[0]));
} else if (op->is_intrinsic(Call::bitwise_and)) {
print_assignment(op->type, print_expr(op->args[0]) + " & " + print_expr(op->args[1]));
} else if (op->is_intrinsic(Call::bitwise_or)) {
print_assignment(op->type, print_expr(op->args[0]) + " | " + print_expr(op->args[1]));
} else if (op->is_intrinsic(Call::bitwise_xor)) {
print_assignment(op->type, print_expr(op->args[0]) + " ^ " + print_expr(op->args[1]));
} else if (op->is_intrinsic(Call::div_round_to_zero)) {
print_assignment(op->type, print_expr(op->args[0]) + " / " + print_expr(op->args[1]));
} else if (op->is_intrinsic(Call::mod_round_to_zero)) {
print_assignment(op->type, print_expr(op->args[0]) + " % " + print_expr(op->args[1]));
} else if (op->is_intrinsic(Call::saturating_cast)) {
Expr e = lower_intrinsic(op);
print_expr(e);
return;
} else {
auto it = builtin.find(op->name);
if (it == builtin.end()) {
user_error << "GLSL: unknown function '" << op->name << "' encountered.\n";
}
ostringstream rhs;
rhs << it->second << "(";
for (size_t i = 0; i < op->args.size(); i++) {
if (i > 0) {
rhs << ", ";
}
rhs << print_expr(op->args[i]);
}
rhs << ")";
print_assignment(op->type, rhs.str());
}
}
void CodeGen_OpenGLCompute_C::visit(const Cast *op) {
Type value_type = op->value.type();
// If both types are represented by the same GLSL type, no explicit cast
// is necessary.
Type target_type = map_type(op->type);
if (target_type == map_type(value_type)) {
Expr value = op->value;
if (value_type.code() == Type::Float) {
// float->int conversions may need explicit truncation if an
// integer type is embedded into a float. (Note: overflows are
// considered undefined behavior, so we do nothing about values
// that are out of range of the target type.)
if (op->type.code() == Type::UInt) {
value = simplify(floor(value));
} else if (op->type.code() == Type::Int) {
value = simplify(trunc(value));
}
}
// FIXME: Overflow is not UB for most Halide types
// https://github.com/halide/Halide/issues/4975
value.accept(this);
} else {
print_assignment(target_type, print_type(target_type) + "(" + print_expr(op->value) + ")");
}
}
void CodeGen_OpenGLCompute_C::visit(const For *loop) {
user_assert(loop->for_type != ForType::GPULane)
<< "The OpenGLCompute backend does not support the gpu_lanes() scheduling directive.";
if (CodeGen_GPU_Dev::is_gpu_var(loop->name)) {
internal_assert((loop->for_type == ForType::GPUBlock) ||
(loop->for_type == ForType::GPUThread))
<< "kernel loop must be either gpu block or gpu thread\n";
internal_assert(is_const_zero(loop->min));
debug(4) << "loop extent is " << loop->extent << "\n";
//
// Need to extract workgroup size.
//
int index = thread_loop_workgroup_index(loop->name);
if (index >= 0) {
const IntImm *int_limit = loop->extent.as<IntImm>();
user_assert(int_limit != nullptr) << "For OpenGLCompute workgroup size must be a constant integer.\n";
int new_workgroup_size = int_limit->value;
user_assert(workgroup_size[index] == 0 ||
workgroup_size[index] == new_workgroup_size)
<< "OpenGLCompute requires all gpu kernels have same workgroup size, "
<< "but two different ones were encountered " << workgroup_size[index]
<< " and " << new_workgroup_size
<< " in dimension " << index << ".\n";
workgroup_size[index] = new_workgroup_size;
debug(4) << "Workgroup size for index " << index << " is " << workgroup_size[index] << "\n";
}
stream << get_indent() << print_type(Int(32)) << " " << print_name(loop->name)
<< " = int(" << simt_intrinsic(loop->name) << ");\n";
loop->body.accept(this);
} else {
user_assert(loop->for_type != ForType::Parallel)
<< "Cannot use parallel loops inside OpenGLCompute kernel\n";
CodeGen_C::visit(loop);
}
}
void CodeGen_OpenGLCompute_C::visit(const Ramp *op) {
if (op->lanes > 4) {
internal_error << "GLSL: ramp lanes " << op->lanes << " is not supported\n";
}
ostringstream rhs;
// Print the sequence vec(0, 1, 2, ...).
rhs << print_type(op->type) << "(";
for (int i = 0; i < op->type.lanes(); i++) {
rhs << i;
if (i != op->type.lanes() - 1) {
rhs << ", ";
}
}
rhs << ")";
// Multiply by the stride and add the base.
rhs << " * " << print_expr(op->stride) << " + " << print_expr(op->base);
print_assignment(op->type, rhs.str());
}
void CodeGen_OpenGLCompute_C::visit(const Broadcast *op) {
string id_value = print_expr(op->value);
ostringstream oss;
oss << print_type(op->type.with_lanes(op->lanes)) << "(" << id_value << ")";
print_assignment(op->type.with_lanes(op->lanes), oss.str());
}
void CodeGen_OpenGLCompute_C::visit(const Load *op) {
user_assert(is_const_one(op->predicate)) << "GLSL: predicated load is not supported.\n";
// https://github.com/halide/Halide/issues/4975
string name = print_name(op->name);
if (!allocations.contains(op->name)) {
name += ".data";
}
// If the index is scalar, just index the buffer using the index.
if (op->type.is_scalar()) {
internal_assert(!buffer_is_vector[op->name]);
string index_id = print_expr(op->index);
string rhs = name + "[" + index_id + "]";
print_assignment(op->type, rhs);
return;
}
// If this is a dense vector load and the buffer has a vector base type,
// then index the buffer using the base of the ramp divided by the number
// of lanes.
Expr ramp_base = strided_ramp_base(op->index);
if (ramp_base.defined() && buffer_is_vector[op->name]) {
string index_id = print_expr(ramp_base / op->type.lanes());
string rhs = name + "[" + index_id + "]";
print_assignment(op->type, rhs);
return;
}
// Gather vector elements.
internal_assert(op->type.is_vector());
internal_assert(!buffer_is_vector[op->name]);
string index_id = print_expr(op->index);
string rhs = print_type(op->type) + "(";
for (int i = 0; i < op->type.lanes(); i++) {
rhs += name + "[" + index_id + "[" + std::to_string(i) + "]]";
if (i != op->type.lanes() - 1) {
rhs += ", ";
}
}
rhs += ")";
print_assignment(op->type, rhs);
}
void CodeGen_OpenGLCompute_C::visit(const Store *op) {
user_assert(is_const_one(op->predicate)) << "GLSL: predicated store is not supported.\n";
// https://github.com/halide/Halide/issues/4975
string name = print_name(op->name);
if (!allocations.contains(op->name)) {
name += ".data";
}
string value_id = print_expr(op->value);
// If the index is scalar, just index the buffer using the index.
if (op->value.type().is_scalar()) {
internal_assert(!buffer_is_vector[op->name]);
string index_id = print_expr(op->index);
stream << get_indent() << name << "[" << index_id << "] = ";
stream << value_id << ";\n";
// Need a cache clear on stores to avoid reusing stale loaded
// values from before the store.
cache.clear();
return;
}
// If this is a dense vector store and the buffer has a vector base type,
// then index the buffer using the base of the ramp divided by the number
// of lanes.
Expr ramp_base = strided_ramp_base(op->index);
if (ramp_base.defined() && buffer_is_vector[op->name]) {
string index_id = print_expr(ramp_base / op->value.type().lanes());
stream << get_indent() << name << "[" << index_id << "] = ";
stream << value_id << ";\n";
// Need a cache clear on stores to avoid reusing stale loaded
// values from before the store.
cache.clear();
return;
}
// Scatter vector elements.
internal_assert(op->value.type().is_vector());
internal_assert(!buffer_is_vector[op->name]);
string index_id = print_expr(op->index);
for (int i = 0; i < op->value.type().lanes(); i++) {
string sub_index_id = index_id + "[" + std::to_string(i) + "]";
stream << get_indent() << name << "[" << sub_index_id << "] = ";
stream << value_id << "[" << std::to_string(i) << "];\n";
}
// Need a cache clear on stores to avoid reusing stale loaded
// values from before the store.
cache.clear();
}
void CodeGen_OpenGLCompute_C::visit(const Select *op) {
ostringstream rhs;
string true_val = print_expr(op->true_value);
string false_val = print_expr(op->false_value);
string cond = print_expr(op->condition);
if (op->type.is_scalar()) {
rhs << cond << " ? " << true_val << " : " << false_val;
} else {
rhs << print_type(op->type) << "(";
for (int i = 0; i < op->type.lanes(); i++) {
string index = "[" + std::to_string(i) + "]";
rhs << cond << index << " ? "
<< true_val << index << " : "
<< false_val << index;
if (i != op->type.lanes() - 1) {
rhs << ", ";
}
}
rhs << ")";
}
print_assignment(op->type, rhs.str());
}
class CodeGen_OpenGLCompute_Dev : public CodeGen_GPU_Dev {
public:
CodeGen_OpenGLCompute_Dev(const Target &target);
// CodeGen_GPU_Dev interface
void add_kernel(Stmt stmt,
const std::string &name,
const std::vector<DeviceArgument> &args) override;
void init_module() override;
std::vector<char> compile_to_src() override;
std::string get_current_kernel_name() override;
void dump() override;
std::string print_gpu_name(const std::string &name) override;
std::string api_unique_name() override {
return "openglcompute";
}
bool kernel_run_takes_types() const override {
return true;
}
protected:
std::ostringstream src_stream;
std::string cur_kernel_name;
CodeGen_OpenGLCompute_C glc;
};
CodeGen_OpenGLCompute_Dev::CodeGen_OpenGLCompute_Dev(const Target &target)
: glc(src_stream, target) {
}
void CodeGen_OpenGLCompute_Dev::add_kernel(Stmt s,
const string &name,
const vector<DeviceArgument> &args) {
debug(2) << "CodeGen_OpenGLCompute_Dev::compile " << name << "\n";
// TODO: do we have to uniquify these names, or can we trust that they are safe?
cur_kernel_name = name;
glc.add_kernel(s, name, args);
}
namespace {
class FindSharedAllocations : public IRVisitor {
using IRVisitor::visit;
void visit(const Allocate *op) override {
op->body.accept(this);
if (op->memory_type == MemoryType::GPUShared) {
allocs.push_back(op);
}
}
public:
vector<const Allocate *> allocs;
};
// Check if all loads and stores to the member 'buffer' are dense, aligned, and
// have the same number of lanes. If this is indeed the case then the 'lanes'
// member stores the number of lanes in those loads and stores.
class CheckAlignedDenseVectorLoadStore : public IRVisitor {
public:
// True if all loads and stores from the buffer are dense, aligned, and all
// have the same number of lanes, false otherwise.
bool are_all_dense = true;
// The number of lanes in the loads and stores. If the number of lanes is
// variable, then are_all_dense is set to false regardless, and this value
// is undefined. Initially set to -1 before any dense operation is
// discovered.
int lanes = -1;
CheckAlignedDenseVectorLoadStore(string buffer)
: buffer(std::move(buffer)) {
}
private:
// The name of the buffer to check.
string buffer;
using IRVisitor::visit;
void visit(const Load *op) override {
IRVisitor::visit(op);
if (op->name != buffer) {
return;
}
if (op->type.is_scalar()) {
are_all_dense = false;
return;
}
Expr ramp_base = strided_ramp_base(op->index);
if (!ramp_base.defined()) {
are_all_dense = false;
return;
}
if ((op->alignment.modulus % op->type.lanes() != 0) ||
(op->alignment.remainder % op->type.lanes() != 0)) {
are_all_dense = false;
return;
}
if (lanes != -1 && op->type.lanes() != lanes) {
are_all_dense = false;
return;
}
lanes = op->type.lanes();
}
void visit(const Store *op) override {
IRVisitor::visit(op);
if (op->name != buffer) {
return;
}
if (op->value.type().is_scalar()) {
are_all_dense = false;
return;
}
Expr ramp_base = strided_ramp_base(op->index);
if (!ramp_base.defined()) {
are_all_dense = false;
return;
}
if ((op->alignment.modulus % op->value.type().lanes() != 0) ||
(op->alignment.remainder % op->value.type().lanes() != 0)) {
are_all_dense = false;
return;
}
if (lanes != -1 && op->value.type().lanes() != lanes) {
are_all_dense = false;
return;
}
lanes = op->value.type().lanes();
}
};
} // namespace
void CodeGen_OpenGLCompute_C::add_kernel(const Stmt &s,
const string &name,
const vector<DeviceArgument> &args) {
debug(2) << "Adding OpenGLCompute kernel " << name << "\n";
cache.clear();
if (target.os == Target::Android) {
stream << "#version 310 es\n"
<< "#extension GL_ANDROID_extension_pack_es31a : require\n";
} else if (target.has_feature(Target::EGL)) {
stream << "#version 310 es\n";
} else {
stream << "#version 430\n";
}
stream << "float float_from_bits(int x) { return intBitsToFloat(int(x)); }\n";
stream << "#define halide_maybe_unused(x) (void)(x)\n";
for (size_t i = 0; i < args.size(); i++) {
if (args[i].is_buffer) {
//
// layout(binding = 10) buffer buffer10 {
// vec3 data[];
// } inBuffer;
//
CheckAlignedDenseVectorLoadStore check_dense(args[i].name);
s.accept(&check_dense);
int lanes = check_dense.are_all_dense ? check_dense.lanes : 1;
buffer_is_vector[args[i].name] = lanes > 1;
stream << "layout(binding=" << i << ")"
<< " buffer buffer" << i << " { "
<< print_type(args[i].type.with_lanes(lanes)) << " data[]; } "
<< print_name(args[i].name) << ";\n";
} else {
stream << "layout(location = " << i << ") uniform " << print_type(args[i].type)
<< " " << print_name(args[i].name) << ";\n";
}
}
// Find all the shared allocations and declare them at global scope.
FindSharedAllocations fsa;
s.accept(&fsa);
for (const Allocate *op : fsa.allocs) {
internal_assert(op->extents.size() == 1 && is_const(op->extents[0]));
stream << "shared "
<< print_type(op->type) << " "
<< print_name(op->name) << "["
<< op->extents[0] << "];\n";
}
// We'll figure out the workgroup size while traversing the stmt
workgroup_size[0] = 0;
workgroup_size[1] = 0;
workgroup_size[2] = 0;
stream << "void main()\n{\n";
indent += 2;
print(s);
indent -= 2;
stream << "}\n";
// Declare the workgroup size.
indent += 2;
stream << "layout(local_size_x = " << workgroup_size[0];
if (workgroup_size[1] > 1) {
stream << ", local_size_y = " << workgroup_size[1];
}
if (workgroup_size[2] > 1) {
stream << ", local_size_z = " << workgroup_size[2];
}
stream << ") in;\n// end of kernel " << name << "\n";
indent -= 2;
}
void CodeGen_OpenGLCompute_Dev::init_module() {
src_stream.str("");
src_stream.clear();
cur_kernel_name = "";
}
void CodeGen_OpenGLCompute_C::visit(const Allocate *op) {
debug(2) << "OpenGLCompute: Allocate " << op->name << " of type " << op->type << " on device\n";
stream << get_indent();
Allocation alloc;
alloc.type = op->type;
allocations.push(op->name, alloc);
internal_assert(!op->extents.empty());
Expr extent = 1;
for (const Expr &e : op->extents) {
extent *= e;
}
extent = simplify(extent);
internal_assert(is_const(extent));
if (op->memory_type != MemoryType::GPUShared) {
stream << "{\n";
indent += 2;
stream << get_indent();
// Shared allocations were already declared at global scope.
stream << print_type(op->type) << " "
<< print_name(op->name) << "["
<< op->extents[0] << "];\n";
}
op->body.accept(this);
if (op->memory_type != MemoryType::GPUShared) {
indent -= 2;
stream << get_indent() << "}\n";
}
buffer_is_vector[op->name] = op->type.is_vector();
}
void CodeGen_OpenGLCompute_C::visit(const Free *op) {
debug(2) << "OpenGLCompute: Free on device for " << op->name << "\n";
allocations.pop(op->name);
}
void CodeGen_OpenGLCompute_C::visit(const Evaluate *op) {
if (is_const(op->value)) {
return;
}
print_expr(op->value);
}
void CodeGen_OpenGLCompute_C::visit(const IntImm *op) {
if (op->type == Int(32)) {
// GL seems to interpret some large int immediates as uints.
id = "int(" + std::to_string(op->value) + ")";
} else {
id = print_type(op->type) + "(" + std::to_string(op->value) + ")";
}
}
vector<char> CodeGen_OpenGLCompute_Dev::compile_to_src() {
string str = src_stream.str();
debug(1) << "GLSL Compute source:\n"
<< str << "\n";
vector<char> buffer(str.begin(), str.end());
buffer.push_back(0);
return buffer;
}
string CodeGen_OpenGLCompute_Dev::get_current_kernel_name() {
return cur_kernel_name;
}
void CodeGen_OpenGLCompute_Dev::dump() {
std::cerr << src_stream.str() << "\n";
}
std::string CodeGen_OpenGLCompute_Dev::print_gpu_name(const std::string &name) {
return name;
}
} // namespace
std::unique_ptr<CodeGen_GPU_Dev> new_CodeGen_OpenGLCompute_Dev(const Target &target) {
return std::make_unique<CodeGen_OpenGLCompute_Dev>(target);
}
} // namespace Internal
} // namespace Halide