Function passesJitHeuristics

src/backend/cuda/Array.cpp:271–350

///    limitation on the platform. For NVIDIA this is 4096 bytes. The
template<typename T>
kJITHeuristics passesJitHeuristics(span<Node *> root_nodes) {
    if (!evalFlag()) { return kJITHeuristics::Pass; }
    static auto getLogger = [&] { return spdlog::get("jit"); };
    for (Node *n : root_nodes) {
        if (n->getHeight() > static_cast<int>(getMaxJitSize())) {
            AF_TRACE(
                "JIT tree evaluated because of tree height exceeds limit: {} > "
                "{}",
                n->getHeight(), getMaxJitSize());
            return kJITHeuristics::TreeHeight;
        }
    }

    // A lightweight check based on the height of the node. This is an
    // inexpensive operation and does not traverse the JIT tree.
    int heightCheckLimit = 6;
    bool atHeightLimit =
        std::any_of(std::begin(root_nodes), std::end(root_nodes),
                    [heightCheckLimit](Node *n) {
                        return (n->getHeight() + 1 >= heightCheckLimit);
                    });
    if (atHeightLimit || getMemoryPressure() >= getMemoryPressureThreshold()) {
        // The size of the parameters without any extra arguments from the
        // JIT tree. This includes one output Param object and 4 integers.
        size_t base_param_size =
            sizeof(Param<T>) * root_nodes.size() + (4 * sizeof(uint));

        // extra padding for safety to avoid failure during compilation
        constexpr size_t jit_padding_size = 256;  //@umar dontfix!
        // This is the maximum size of the params that can be allowed by the
        // CUDA platform.
        size_t max_param_size = 4096 - base_param_size - jit_padding_size;

        struct tree_info {
            size_t total_buffer_size;
            size_t num_buffers;
            size_t param_scalar_size;
        };
        NodeIterator<> end_node;
        tree_info info = tree_info{0, 0, 0};

        for (Node *n : root_nodes) {
            info = accumulate(
                NodeIterator<>(n), end_node, info,
                [](tree_info &prev, const Node &node) {
                    if (node.isBuffer()) {
                        const auto &buf_node =
                            static_cast<const BufferNode<T> &>(node);
                        // getBytes returns the size of the data Array.
                        // Sub arrays will be represented by their
                        // parent size.
                        prev.total_buffer_size += buf_node.getBytes();
                        prev.num_buffers++;
                    } else {
                        prev.param_scalar_size += node.getParamBytes();
                    }
                    return prev;
                });