arrayfire · 9prady9 · Feb 22, 2021 · Jan 1, 2021 · Dec 30, 2020 · Jan 4, 2021
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -1,4 +1,4 @@
-# Copyright (c) 2020, ArrayFire
+# Copyright (c) 2021, ArrayFire
 # All rights reserved.
 #
 # This file is distributed under 3-clause BSD license.
@@ -7,7 +7,7 @@
 
 cmake_minimum_required(VERSION 3.5)
 
-project(ArrayFire VERSION 3.8.0 LANGUAGES C CXX)
+project(ArrayFire VERSION 3.9.0 LANGUAGES C CXX)
 
 set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} "${CMAKE_CURRENT_SOURCE_DIR}/CMakeModules")
 
@@ -107,15 +107,19 @@ mark_as_advanced(
   AF_BUILD_FRAMEWORK
   AF_INSTALL_STANDALONE
   AF_WITH_CPUID
+  Boost_INCLUDE_DIR
   CUDA_HOST_COMPILER
+  CUDA_SDK_ROOT_DIR
   CUDA_USE_STATIC_CUDA_RUNTIME
   CUDA_rt_LIBRARY
   SPDLOG_BUILD_EXAMPLES
   SPDLOG_BUILD_TESTING
   ADDR2LINE_PROGRAM
   Backtrace_LIBRARY
   AF_WITH_STATIC_MKL
+  GIT
   )
+mark_as_advanced(CLEAR CUDA_VERSION)
 
 #Configure forge submodule
 #forge is included in ALL target if AF_BUILD_FORGE is ON

diff --git a/CMakeModules/FindMKL.cmake b/CMakeModules/FindMKL.cmake
@@ -265,8 +265,8 @@ function(find_mkl_library)
         if (CMAKE_VERSION VERSION_GREATER 3.14)
           message(VERBOSE "MKL_${mkl_args_NAME}_STATIC_LINK_LIBRARY: ${MKL_${mkl_args_NAME}_STATIC_LINK_LIBRARY}")
         endif()
-        mark_as_advanced(MKL_${mkl_args_NAME}_STATIC_LINK_LIBRARY)
       endif()
+      mark_as_advanced(MKL_${mkl_args_NAME}_STATIC_LINK_LIBRARY)
     endif()
 
     set_target_properties(MKL::${mkl_args_NAME}

diff --git a/CMakeModules/FindcuDNN.cmake b/CMakeModules/FindcuDNN.cmake
@@ -151,6 +151,7 @@ if(cuDNN_INCLUDE_DIRS)
         ${CMAKE_INSTALL_PREFIX}
       PATH_SUFFIXES lib lib64 bin lib/x64 bin/x64
       DOC "cudnn${cudnn_lib_name_infix} link library." )
+    mark_as_advanced(cuDNN${LIB_INFIX}_LINK_LIBRARY)
 
     if(WIN32 AND cuDNN_LINK_LIBRARY)
       find_file(cuDNN${LIB_INFIX}_DLL_LIBRARY

diff --git a/src/api/c/CMakeLists.txt b/src/api/c/CMakeLists.txt
@@ -105,6 +105,8 @@ target_sources(c_api_interface
     ${CMAKE_CURRENT_SOURCE_DIR}/index.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/internal.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/inverse.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/jit_test_api.h
+    ${CMAKE_CURRENT_SOURCE_DIR}/jit_test_api.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/join.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/lu.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/match_template.cpp

diff --git a/src/api/c/jit_test_api.cpp b/src/api/c/jit_test_api.cpp
@@ -0,0 +1,28 @@
+/*******************************************************
+ * Copyright (c) 2021, ArrayFire
+ * All rights reserved.
+ *
+ * This file is distributed under 3-clause BSD license.
+ * The complete license agreement can be obtained at:
+ * http://arrayfire.com/licenses/BSD-3-Clause
+ ********************************************************/
+
+#include <jit_test_api.h>
+
+#include <backend.hpp>
+#include <common/err_common.hpp>
+#include <platform.hpp>
+
+af_err af_get_max_jit_len(int *jitLen) {
+    *jitLen = detail::getMaxJitSize();
+    return AF_SUCCESS;
+}
+
+af_err af_set_max_jit_len(const int maxJitLen) {
+    try {
+        ARG_ASSERT(1, maxJitLen > 0);
+        detail::getMaxJitSize() = maxJitLen;
+    }
+    CATCHALL;
+    return AF_SUCCESS;
+}
diff --git a/src/api/c/jit_test_api.h b/src/api/c/jit_test_api.h
@@ -0,0 +1,51 @@
+/*******************************************************
+ * Copyright (c) 2021, ArrayFire
+ * All rights reserved.
+ *
+ * This file is distributed under 3-clause BSD license.
+ * The complete license agreement can be obtained at:
+ * http://arrayfire.com/licenses/BSD-3-Clause
+ ********************************************************/
+
+#pragma once
+
+#include <af/defines.h>
+
+#ifdef __cplusplus
+namespace af {
+/// Get the maximum jit tree length for active backend
+///
+/// \returns the maximum length of jit tree from root to any leaf
+AFAPI int getMaxJitLen(void);
+
+/// Set the maximum jit tree length for active backend
+///
+/// \param[in] jit_len is the maximum length of jit tree from root to any
+/// leaf
+AFAPI void setMaxJitLen(const int jitLen);
+}  // namespace af
+#endif  //__cplusplus
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/// Get the maximum jit tree length for active backend
+///
+/// \param[out] jit_len is the maximum length of jit tree from root to any
+/// leaf
+///
+/// \returns Always returns AF_SUCCESS
+AFAPI af_err af_get_max_jit_len(int *jit_len);
+
+/// Set the maximum jit tree length for active backend
+///
+/// \param[in] jit_len is the maximum length of jit tree from root to any
+/// leaf
+///
+/// \returns Always returns AF_SUCCESS
+AFAPI af_err af_set_max_jit_len(const int jit_len);
+
+#ifdef __cplusplus
+}
+#endif
diff --git a/src/api/cpp/CMakeLists.txt b/src/api/cpp/CMakeLists.txt
@@ -45,6 +45,7 @@ target_sources(cpp_api_interface
     ${CMAKE_CURRENT_SOURCE_DIR}/imageio.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/index.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/internal.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/jit_test_api.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/lapack.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/matchTemplate.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/mean.cpp

diff --git a/src/api/cpp/jit_test_api.cpp b/src/api/cpp/jit_test_api.cpp
@@ -0,0 +1,21 @@
+/*******************************************************
+ * Copyright (c) 2021, ArrayFire
+ * All rights reserved.
+ *
+ * This file is distributed under 3-clause BSD license.
+ * The complete license agreement can be obtained at:
+ * http://arrayfire.com/licenses/BSD-3-Clause
+ ********************************************************/
+
+#include <jit_test_api.h>
+#include "error.hpp"
+
+namespace af {
+int getMaxJitLen(void) {
+    int retVal = 0;
+    AF_THROW(af_get_max_jit_len(&retVal));
+    return retVal;
+}
+
+void setMaxJitLen(const int jitLen) { AF_THROW(af_set_max_jit_len(jitLen)); }
+}  // namespace af
diff --git a/src/api/unified/CMakeLists.txt b/src/api/unified/CMakeLists.txt
@@ -22,6 +22,7 @@ target_sources(af
     ${CMAKE_CURRENT_SOURCE_DIR}/image.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/index.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/internal.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/jit_test_api.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/lapack.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/memory.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/ml.cpp

diff --git a/src/api/unified/jit_test_api.cpp b/src/api/unified/jit_test_api.cpp
@@ -0,0 +1,18 @@
+/*******************************************************
+ * Copyright (c) 2021, ArrayFire
+ * All rights reserved.
+ *
+ * This file is distributed under 3-clause BSD license.
+ * The complete license agreement can be obtained at:
+ * http://arrayfire.com/licenses/BSD-3-Clause
+ ********************************************************/
+
+#include <jit_test_api.h>
+
+#include "symbol_manager.hpp"
+
+af_err af_get_max_jit_len(int *jitLen) { CALL(af_get_max_jit_len, jitLen); }
+
+af_err af_set_max_jit_len(const int jitLen) {
+    CALL(af_set_max_jit_len, jitLen);
+}
diff --git a/src/backend/common/jit/NaryNode.hpp b/src/backend/common/jit/NaryNode.hpp
@@ -98,8 +98,7 @@ common::Node_ptr createNaryNode(
 
     common::Node_ptr ptr = createNode(childNodes);
 
-    switch (static_cast<kJITHeuristics>(
-        detail::passesJitHeuristics<Ti>(ptr.get()))) {
+    switch (detail::passesJitHeuristics<Ti>(ptr.get())) {
         case kJITHeuristics::Pass: {
             return ptr;
         }
@@ -113,7 +112,6 @@ common::Node_ptr createNaryNode(
                     max_height       = childNodes[i]->getHeight();
                 }
             }
-
             children[max_height_index]->eval();
             return createNaryNode<Ti, N>(odims, createNode, move(children));
         }

diff --git a/src/backend/cpu/Array.cpp b/src/backend/cpu/Array.cpp
@@ -248,7 +248,7 @@ Array<T> createEmptyArray(const dim4 &dims) {
 template<typename T>
 kJITHeuristics passesJitHeuristics(Node *root_node) {
     if (!evalFlag()) { return kJITHeuristics::Pass; }
-    if (root_node->getHeight() >= static_cast<int>(getMaxJitSize())) {
+    if (root_node->getHeight() > static_cast<int>(getMaxJitSize())) {
         return kJITHeuristics::TreeHeight;
     }
 

diff --git a/src/backend/cpu/platform.cpp b/src/backend/cpu/platform.cpp
@@ -104,14 +104,14 @@ void devprop(char* d_name, char* d_platform, char* d_toolkit, char* d_compute) {
     snprintf(d_compute, 10, "%s", "0.0");
 }
 
-unsigned getMaxJitSize() {
-    const int MAX_JIT_LEN = 100;
-
-    thread_local int length = 0;
-    if (length == 0) {
+int& getMaxJitSize() {
+    constexpr int MAX_JIT_LEN = 100;
+    thread_local int length   = 0;
+    if (length <= 0) {
         string env_var = getEnvVar("AF_CPU_MAX_JIT_LEN");
         if (!env_var.empty()) {
-            length = stoi(env_var);
+            int input_len = std::stoi(env_var);
+            length        = input_len > 0 ? input_len : MAX_JIT_LEN;
         } else {
             length = MAX_JIT_LEN;
         }

diff --git a/src/backend/cpu/platform.hpp b/src/backend/cpu/platform.hpp
@@ -36,7 +36,7 @@ bool isHalfSupported(int device);
 
 void devprop(char* d_name, char* d_platform, char* d_toolkit, char* d_compute);
 
-unsigned getMaxJitSize();
+int& getMaxJitSize();
 
 int getDeviceCount();
 

diff --git a/src/backend/cuda/Array.cpp b/src/backend/cuda/Array.cpp
@@ -254,7 +254,7 @@ Node_ptr Array<T>::getNode() const {
 template<typename T>
 kJITHeuristics passesJitHeuristics(Node *root_node) {
     if (!evalFlag()) { return kJITHeuristics::Pass; }
-    if (root_node->getHeight() >= static_cast<int>(getMaxJitSize())) {
+    if (root_node->getHeight() > static_cast<int>(getMaxJitSize())) {
         return kJITHeuristics::TreeHeight;
     }
 

diff --git a/src/backend/cuda/jit.cpp b/src/backend/cuda/jit.cpp
@@ -205,10 +205,13 @@ static CUfunction getKernel(const vector<Node *> &output_nodes,
 template<typename T>
 void evalNodes(vector<Param<T>> &outputs, const vector<Node *> &output_nodes) {
     size_t num_outputs = outputs.size();
-    int device         = getActiveDeviceId();
-
     if (num_outputs == 0) { return; }
 
+    int device         = getActiveDeviceId();
+    dim_t *outDims     = outputs[0].dims;
+    size_t numOutElems = outDims[0] * outDims[1] * outDims[2] * outDims[3];
+    if (numOutElems == 0) { return; }
+
     // Use thread local to reuse the memory every time you are here.
     thread_local Node_map_t nodes;
     thread_local vector<Node *> full_nodes;
@@ -229,9 +232,7 @@ void evalNodes(vector<Param<T>> &outputs, const vector<Node *> &output_nodes) {
     }
 
     bool is_linear = true;
-    for (auto node : full_nodes) {
-        is_linear &= node->isLinear(outputs[0].dims);
-    }
+    for (auto node : full_nodes) { is_linear &= node->isLinear(outDims); }
 
     CUfunction ker =
         getKernel(output_nodes, output_ids, full_nodes, full_ids, is_linear);
@@ -246,7 +247,7 @@ void evalNodes(vector<Param<T>> &outputs, const vector<Node *> &output_nodes) {
 
     int num_odims = 4;
     while (num_odims >= 1) {
-        if (outputs[0].dims[num_odims - 1] == 1) {
+        if (outDims[num_odims - 1] == 1) {
             num_odims--;
         } else {
             break;
@@ -257,21 +258,20 @@ void evalNodes(vector<Param<T>> &outputs, const vector<Node *> &output_nodes) {
         threads_x = 256;
         threads_y = 1;
 
-        blocks_x_total = divup((outputs[0].dims[0] * outputs[0].dims[1] *
-                                outputs[0].dims[2] * outputs[0].dims[3]),
-                               threads_x);
+        blocks_x_total = divup(
+            (outDims[0] * outDims[1] * outDims[2] * outDims[3]), threads_x);
 
         int repeat_x = divup(blocks_x_total, max_blocks_x);
         blocks_x     = divup(blocks_x_total, repeat_x);
     } else {
         threads_x = 32;
         threads_y = 8;
 
-        blocks_x_ = divup(outputs[0].dims[0], threads_x);
-        blocks_y_ = divup(outputs[0].dims[1], threads_y);
+        blocks_x_ = divup(outDims[0], threads_x);
+        blocks_y_ = divup(outDims[1], threads_y);
 
-        blocks_x = blocks_x_ * outputs[0].dims[2];
-        blocks_y = blocks_y_ * outputs[0].dims[3];
+        blocks_x = blocks_x_ * outDims[2];
+        blocks_y = blocks_y_ * outDims[3];
 
         blocks_z = divup(blocks_y, max_blocks_y);
         blocks_y = divup(blocks_y, blocks_z);

diff --git a/src/backend/cuda/platform.cpp b/src/backend/cuda/platform.cpp
@@ -325,14 +325,14 @@ string getCUDARuntimeVersion() noexcept {
     }
 }
 
-unsigned getMaxJitSize() {
-    const int MAX_JIT_LEN = 100;
-
-    thread_local int length = 0;
-    if (length == 0) {
+int &getMaxJitSize() {
+    constexpr int MAX_JIT_LEN = 100;
+    thread_local int length   = 0;
+    if (length <= 0) {
         std::string env_var = getEnvVar("AF_CUDA_MAX_JIT_LEN");
         if (!env_var.empty()) {
-            length = std::stoi(env_var);
+            int input_len = std::stoi(env_var);
+            length        = input_len > 0 ? input_len : MAX_JIT_LEN;
         } else {
             length = MAX_JIT_LEN;
         }

diff --git a/src/backend/cuda/platform.hpp b/src/backend/cuda/platform.hpp
@@ -76,7 +76,7 @@ bool isHalfSupported(int device);
 
 void devprop(char* d_name, char* d_platform, char* d_toolkit, char* d_compute);
 
-unsigned getMaxJitSize();
+int& getMaxJitSize();
 
 int getDeviceCount();