#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