Function passesJitHeuristics

src/backend/opencl/Array.cpp:327–423

///    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 common::loggerFactory("jit"); };
    for (const 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;
        }
    }

    bool isBufferLimit = getMemoryPressure() >= getMemoryPressureThreshold();
    auto platform      = getActivePlatformVendor();

    // The Apple platform can have the nvidia card or the AMD card
    bool isIntel = platform == AFCL_PLATFORM_INTEL;

    /// Intels param_size limit is much smaller than the other platforms
    /// so we need to start checking earlier with smaller trees
    int heightCheckLimit =
        isIntel && getDeviceType() == CL_DEVICE_TYPE_GPU ? 3 : 6;

    // A lightweight check based on the height of the node. This is
    // an inexpensive operation and does not traverse the JIT tree.
    bool atHeightLimit =
        std::any_of(std::begin(root_nodes), std::end(root_nodes),
                    [heightCheckLimit](Node *n) {
                        return (n->getHeight() + 1 >= heightCheckLimit);
                    });

    if (atHeightLimit || isBufferLimit) {
        // This is the base parameter size if the kernel had no
        // arguments
        size_t base_param_size =
            (sizeof(T *) + sizeof(KParam)) * root_nodes.size() +
            (3 * sizeof(uint));

        const cl::Device &device = getDevice();
        // typical values:
        //   NVIDIA     = 4096
        //   AMD        = 3520  (AMD A10 iGPU = 1024)
        //   Intel iGPU = 1024
        //
        // Setting the maximum to 5120 bytes to keep the compile times
        // resonable. This still results in large kernels but its not excessive.
        size_t max_param_size =
            min(static_cast<cl::size_type>(5120),
                device.getInfo<CL_DEVICE_MAX_PARAMETER_SIZE>());
        max_param_size -= base_param_size;

        struct tree_info {
            size_t total_buffer_size;
            size_t num_buffers;
            size_t param_scalar_size;
        };