Clean up cuda build when functrace is enabled (deeplearning4j#10197)

* Clean up cuda build when functrace is enabled

* Fix cuda tad cache usage

Author: agibsonccc

libnd4j/CMakeLists.txt

@@ -113,8 +113,10 @@ if (SD_CUDA)
         message("Jetson nano cublas library is ${CUDA_cublas_LIBRARY} and CuSolver library ${CUDA_cusolver_LIBRARY}")
     endif()
 
-
-    set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS}   -allow-unsupported-compiler --ptxas-options=-v  -Xptxas -O1")
+    if("${SD_PTXAS}" STREQUAL  "ON")
+        set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --ptxas-options=-v")
+    endif()
+    set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS}   -allow-unsupported-compiler  ")
     if(SD_KEEP_NVCC_OUTPUT)
         set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS}  --keep ")
     endif()
@@ -157,13 +159,7 @@ foreach(TEMPLATE_FILE ${CUDA_TEMPLATE_FILES})
 endforeach()
 
 
-# Define a function to generate individual CUDA files for each type combination
-# Find and replace the existing function:
-function(genSingleFunctionCuda TEMPLATE_FILE COMBINATION OUTPUT_DIR)
-    # ...existing implementation...
-endfunction()
 
-# REPLACE WITH THIS IMPLEMENTATION:
 
 function(genSingleFunctionCuda TEMPLATE_FILE COMBINATION OUTPUT_DIR)
     # Split the COMBINATION string into a list
@@ -573,8 +569,8 @@ endif()
 # -fsanitize=leak
 if (SD_ANDROID_BUILD)
     set_property(GLOBAL PROPERTY JOB_POOLS one_job=1 two_jobs=2)
-    set(CMAKE_CXX_FLAGS_RELEASE  "${CMAKE_CXX_FLAGS_RELEASE} -O${SD_OPTIMIZATION_LEVEL} -fPIC -Wno-braced-scalar-init -Wno-delete-non-virtual-dtor -Wno-unused-command-line-argument -Wno-dangling-else -D_RELEASE=true")
-    set(CMAKE_CXX_FLAGS_DEBUG  "${CMAKE_CXX_FLAGS_DEBUG} -O${SD_OPTIMIZATION_LEVEL} -g -fPIC -Wno-braced-scalar-init -Wno-delete-non-virtual-dtor -Wno-unused-command-line-argument -Wno-dangling-else")
+    set(CMAKE_CXX_FLAGS_RELEASE  "${CMAKE_CXX_FLAGS_RELEASE} -O${SD_OPTIMIZATION_LEVEL} -fPIC -Wno-return-type -Wno-unknown-pragmas -Wno-braced-scalar-init -Wno-delete-non-virtual-dtor -Wno-unused-command-line-argument -Wno-dangling-else -D_RELEASE=true")
+    set(CMAKE_CXX_FLAGS_DEBUG  "${CMAKE_CXX_FLAGS_DEBUG} -O${SD_OPTIMIZATION_LEVEL} -g -fPIC -Wno-return-type -Wno-unknown-pragmas -Wno-braced-scalar-init -Wno-delete-non-virtual-dtor -Wno-unused-command-line-argument -Wno-dangling-else")
 elseif (APPLE)
     if(${CMAKE_SYSTEM_PROCESSOR} MATCHES "arm64*" OR "${SD_ARCH}" MATCHES "armv8-a")
         set(SD_ARCH armv8-a)
@@ -583,8 +579,8 @@ elseif (APPLE)
     endif()
 
 
-    set(CMAKE_CXX_FLAGS_RELEASE  "-O${SD_OPTIMIZATION_LEVEL} -fPIC -Wno-braced-scalar-init -Wno-delete-non-virtual-dtor -Wno-unused-command-line-argument -Wno-dangling-else -D__APPLE_OS__=true -D_RELEASE=true")
-    set(CMAKE_CXX_FLAGS_DEBUG  " -O${SD_OPTIMIZATION_LEVEL} -g -fPIC -Wno-braced-scalar-init -Wno-delete-non-virtual-dtor -Wno-unused-command-line-argument -Wno-dangling-else -D__APPLE_OS__=true")
+    set(CMAKE_CXX_FLAGS_RELEASE  "-O${SD_OPTIMIZATION_LEVEL} -fPIC -Wno-return-type -Wno-braced-scalar-init -Wno-unknown-pragmas -Wno-delete-non-virtual-dtor -Wno-unused-command-line-argument -Wno-dangling-else -D__APPLE_OS__=true -D_RELEASE=true")
+    set(CMAKE_CXX_FLAGS_DEBUG  " -O${SD_OPTIMIZATION_LEVEL} -g -fPIC -Wno-return-type -Wno-braced-scalar-init -Wno-unknown-pragmas -Wno-delete-non-virtual-dtor -Wno-unused-command-line-argument -Wno-dangling-else -D__APPLE_OS__=true")
 elseif(WIN32)
     set(SD_X86_BUILD true)
     if (SD_CUDA)

libnd4j/blas/CMakeLists.txt

@@ -200,7 +200,7 @@ elseif(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND NOT CMAKE_SYSTEM_NAME MATCHES "A
         endif()
 
         # Set C++ compiler and flags
-        set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}  -fstack-protector -fstack-protector-all  -Wall -Wextra -Werror -Wno-error=int-in-bool-context -Wno-unused-variable -Wno-error=implicit-fallthrough -Wno-return-type -Wno-unused-parameter -Wno-error=unknown-pragmas -ggdb3 -lpthread -pthread -MT -Bsymbolic -lbfd -rdynamic -lunwind -ldw -ldl -fno-omit-frame-pointer -fno-optimize-sibling-calls -rdynamic -finstrument-functions  -O0 -fPIC")
+        set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}  -fstack-protector -fstack-protector-all  -Wall  -Wextra -Werror -Wno-return-type -Wno-error=int-in-bool-context -Wno-unused-variable -Wno-error=implicit-fallthrough -Wno-return-type -Wno-unused-parameter -Wno-error=unknown-pragmas -ggdb3 -lpthread -pthread -MT -Bsymbolic -lbfd -rdynamic -lunwind -ldw -ldl -fno-omit-frame-pointer -fno-optimize-sibling-calls -rdynamic -finstrument-functions  -O0 -fPIC")
         add_compile_definitions(SD_GCC_FUNCTRACE)
     endif()
 endif()
@@ -262,7 +262,7 @@ if(SD_CUDA)
 
         if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
             if(SD_GCC_FUNCTRACE STREQUAL "ON")
-                set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}  -Werror -Wall   -Wno-unused-variable -Wno-unused-parameter  -Wreturn-type -W -ggdb3 -fPIC -DSD_GCC_FUNCTRACE=1 -Bsymbolic -lbfd -rdynamic -lunwind -ldw -ldl -fno-omit-frame-pointer -fno-optimize-sibling-calls -finstrument-functions  -O0")
+                set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}  -Werror -Wall   -Wno-return-type  -Wno-unknown-pragmas  -Wno-unused-variable -Wno-unused-parameter  -Wreturn-type -W -ggdb3 -fPIC -DSD_GCC_FUNCTRACE=1 -Bsymbolic -lbfd -rdynamic -lunwind -ldw -ldl -fno-omit-frame-pointer -fno-optimize-sibling-calls -finstrument-functions  -O0")
                 set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler=-fPIC --device-debug -lineinfo -G")
                 add_compile_definitions(SD_GCC_FUNCTRACE)
             else()
@@ -306,7 +306,7 @@ if(SD_CUDA)
         endif()
 
         # Cap the number of registers to prevent resource exhaustion
-        set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --maxrregcount=40 ")
+        set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -maxrregcount=128  ")
 
         # Define CUDA Architectures
         string(TOLOWER "${COMPUTE}" COMPUTE_CMP)
@@ -371,6 +371,11 @@ if(SD_CUDA)
         endif()
        include(${CMAKE_CURRENT_SOURCE_DIR}/../cmake/TypeMST.cmake)
 
+        # 2. Process the compilation_units templates
+        file(GLOB CUDA_COMPILATION_UNITS ../include/loops/cuda/compilation_units/*.cu.in)
+        foreach(FL_ITEM ${CUDA_COMPILATION_UNITS})
+            genCompilation(${FL_ITEM})
+        endforeach()
 
         # Decide whether to use all combinations or optimized MST combinations
        set(SD_USE_MST_TYPES ON)
@@ -423,7 +428,7 @@ if(SD_CUDA)
         file(GLOB_RECURSE INDEXING_SOURCES ../include/indexing/*.cpp ../include/indexing/*.h)
         file(GLOB_RECURSE LOOPS_SOURCES ../include/loops/impl/*.cpp ../include/loops/*.h ../include/loops/*.chpp)
         file(GLOB_RECURSE LEGACY_SOURCES ../include/legacy/impl/*.cpp ../include/legacy/*.cu ../include/legacy/*.h)
-        file(GLOB_RECURSE LOOPS_SOURCES_CUDA ../include/loops/*.cu)
+        file(GLOB_RECURSE LOOPS_SOURCES_CUDA ../include/loops/*.cu ../include/loops/cuda/**/*.cu)
         file(GLOB_RECURSE COMPILATION_UNITS ../include/loops/cuda/compilation_units/*.cu.in ../include/ops/impl/compilation_units/*.cpp.in)
         file(GLOB_RECURSE COMPILATION_UNITS ../include/loops/cuda/compilation_units/*.cu.in ../include/loops/cuda/comb_compilation_units/*.cu.in  ../include/ops/impl/compilation_units/*.cpp.in)
 
@@ -485,6 +490,8 @@ if(SD_CUDA)
                     ${CUSTOMOPS_ONEDNN_SOURCES}
                     ${CUSTOMOPS_ARMCOMPUTE_SOURCES}
                     ${CUSTOMOPS_GENERIC_SOURCES}
+                    ${LOOPS_SOURCES_CUDA}
+
             )
         else()
             add_library(samediff_obj OBJECT
@@ -509,6 +516,8 @@ if(SD_CUDA)
                     ${CUSTOMOPS_ONEDNN_SOURCES}
                     ${CUSTOMOPS_ARMCOMPUTE_SOURCES}
                     ${CUSTOMOPS_GENERIC_SOURCES}
+                    ${LOOPS_SOURCES_CUDA}
+
             )
         endif()
 

libnd4j/buildnativeoperations.sh

@@ -94,7 +94,7 @@ PRINT_MATH="OFF"
 KEEP_NVCC="OFF"
 PREPROCESS="ON"  # Initialize PREPROCESS variable
 CMAKE_ARGUMENTS=""
-
+PTXAS_INFO="OFF"
 while [[ $# -gt 0 ]]
 do
     key="$1"
@@ -106,6 +106,10 @@ do
                 CMAKE_ARGUMENTS="$CMAKE_ARGUMENTS -DGENERATE_FLATC=ON"
                 shift # past argument
                 ;;
+    -ptxas|--ptxas-info)
+            PTXAS_INFO="$value"
+            shift # past argument
+            ;;
         -ol|--optimization-level)
             OPTIMIZATION_LEVEL="$value"
             shift # past argument
@@ -712,6 +716,7 @@ if [ "$LOG_OUTPUT" == "none" ]; then
         -DPRINT_INDICES="$PRINT_INDICES" \
         -DSD_KEEP_NVCC_OUTPUT="$KEEP_NVCC" \
         -DSD_GCC_FUNCTRACE="$FUNC_TRACE" \
+        -DSD_PTXAS="$PTXAS_INFO" \
         "$BLAS_ARG" \
         "$ARCH_ARG" \
         "$NAME_ARG" \
@@ -739,6 +744,7 @@ else
         -DPRINT_INDICES="$PRINT_INDICES" \
         -DSD_KEEP_NVCC_OUTPUT="$KEEP_NVCC" \
         -DSD_GCC_FUNCTRACE="$FUNC_TRACE" \
+        -DSD_PTXAS="$PTXAS_INFO" \
         "$BLAS_ARG" \
         "$ARCH_ARG" \
         "$NAME_ARG" \
@@ -773,6 +779,7 @@ if [ "$PREPROCESS" == "ON" ]; then
            -DSD_KEEP_NVCC_OUTPUT="$KEEP_NVCC" \
            -DSD_GCC_FUNCTRACE="$FUNC_TRACE" \
             -DSD_PREPROCESS="$PREPROCESS" \
+           -DSD_PTXAS="$PTXAS_INFO" \
            "$BLAS_ARG" \
            "$ARCH_ARG" \
            "$NAME_ARG" \
@@ -804,6 +811,7 @@ if [ "$PREPROCESS" == "ON" ]; then
            "$ARCH_ARG" \
            "$NAME_ARG" \
            "$OP_OUTPUT_FILE_ARG" \
+           -DSD_PTXAS="$PTXAS_INFO" \
            -DSD_SANITIZE="${SANITIZE}" \
            -DSD_CHECK_VECTORIZATION="${CHECK_VECTORIZATION}" \
            "$USE_LTO" \

libnd4j/include/array/NDArray.h

@@ -2003,7 +2003,7 @@ void * _bufferWithOffset(LongType offset,DataBuffer *buffer) {
 //note this is meant to be used with primary() (host side/cpu) use specialBuffer() for device side buffers
 void *NDArray::buffer() {
   BUILD_SINGLE_SELECTOR(dataType(), return _bufferWithOffset, (offset(),getDataBuffer()),SD_COMMON_TYPES);
-
+  return nullptr;
 }
 
 //////////////////////////////////////////////////////////////////////////

libnd4j/include/array/NDArray.hXX

@@ -2639,7 +2639,7 @@ void NDArray::copyDataForAssign(NDArray& thisArray,  NDArray& other, const sd::L
                                           other.buffer(),
                                           otherShapeInfo,
                                           other.specialBuffer(),
-                                          otherShapeInfo,
+                                          other.specialShapeInfo(),
                                           nullptr
         , allowParallelism);
     thisArray.registerSpecialUse({&thisArray}, {&other}); // Registering through the instance

libnd4j/include/array/TadDescriptor.h

@@ -28,7 +28,7 @@
 namespace sd {
 class SD_LIB_EXPORT TadDescriptor {
  private:
-  ShapeDescriptor _originalShape;
+  sd::LongType *_originalShape;
 
   std::vector<LongType> _axis;
 
@@ -37,8 +37,6 @@ class SD_LIB_EXPORT TadDescriptor {
  public:
   explicit TadDescriptor(const LongType *originalShape, const LongType *dimensions, const LongType length,
                          const bool keepUnitiesInShape = false);
-  explicit TadDescriptor(const ShapeDescriptor &descriptor, const std::vector<LongType> &dimensions,
-                         const bool keepUnitiesInShape = false);
   ~TadDescriptor() = default;
 
 
@@ -56,8 +54,7 @@ class SD_LIB_EXPORT TadDescriptor {
   bool operator<(const TadDescriptor &other) const;
 
   std::vector<LongType> &axis();
-  ShapeDescriptor &originalShape();
-  ShapeDescriptor const &originalShapeConst() const;
+  LongType *originalShape();
   bool areUnitiesinShape() const;
 };
 }  // namespace sd

libnd4j/include/array/cuda/CudaPointerDeallocator.cu

@@ -25,6 +25,23 @@
 
 namespace sd {
 
-void CudaPointerDeallocator::release(void *ptr) { cudaFree(ptr); }
+void CudaPointerDeallocator::release(void *ptr) {
+  if (ptr == nullptr) return;
 
+  // Check if this is a valid device pointer before freeing
+  cudaPointerAttributes attributes;
+  cudaError_t result = cudaPointerGetAttributes(&attributes, ptr);
+
+  if (result == cudaSuccess) {
+    // Only free if it's a regular device pointer
+    // cudaMemoryTypeDevice is for regular allocations we can free
+    if (attributes.type == cudaMemoryTypeDevice) {
+      cudaFree(ptr);
+    }
+    // Don't free other types (like constant memory)
+  } else {
+    // Clear the error and don't try to free this pointer
+    cudaGetLastError(); // Clear the error state
+  }
+}
 }  // namespace sd

libnd4j/include/array/impl/ExtraArguments.cpp

@@ -79,7 +79,7 @@ void ExtraArguments::convertAndCopy(Pointer pointer, LongType offset) {
 #ifdef __CUDABLAS__
   // TODO: maybe make it asynchronous eventually?
   cudaMemcpy(pointer, target, length * DataTypeUtils::sizeOf(DataTypeUtils::fromT<T>()), cudaMemcpyHostToDevice);
-  delete target;
+  delete[] target;
 #endif
 }
 BUILD_SINGLE_TEMPLATE(template SD_LIB_EXPORT void ExtraArguments::convertAndCopy,

libnd4j/include/array/impl/TadDescriptor.cpp

@@ -33,7 +33,6 @@ TadDescriptor::TadDescriptor(const TadDescriptor &other) {
 
 TadDescriptor::TadDescriptor(const LongType *originalShape, const LongType *dimensions, const LongType length,
                              const bool keepUnitiesInShape) {
-  ShapeDescriptor *descriptor = new ShapeDescriptor(originalShape, false);
 
   _axis.resize(length);
   for (LongType e = 0; e < length; e++) {
@@ -42,19 +41,10 @@ TadDescriptor::TadDescriptor(const LongType *originalShape, const LongType *dime
 
   if (length > 1) std::sort(_axis.begin(), _axis.end());
 
-  _originalShape = *descriptor;
+  _originalShape = const_cast<sd::LongType *>(originalShape);
   _unitiesInShape = keepUnitiesInShape;
 }
 
-TadDescriptor::TadDescriptor(const ShapeDescriptor &descriptor, const std::vector<LongType> &dimensions,
-                             const bool keepUnitiesInShape) {
-  _originalShape = descriptor;
-  _axis = dimensions;
-  _unitiesInShape = keepUnitiesInShape;
-
-  if (_axis.size() > 1) std::sort(_axis.begin(), _axis.end());
-}
-
 bool TadDescriptor::operator==(const TadDescriptor &other) const {
   return std::tie(_originalShape, _axis, _unitiesInShape) ==
          std::tie(other._originalShape, other._axis, other._unitiesInShape);
@@ -67,9 +57,7 @@ bool TadDescriptor::operator<(const TadDescriptor &other) const {
 
 std::vector<LongType> &TadDescriptor::axis() { return _axis; }
 
-ShapeDescriptor &TadDescriptor::originalShape() { return _originalShape; }
-
-ShapeDescriptor const &TadDescriptor::originalShapeConst() const { return _originalShape; }
+LongType *TadDescriptor::originalShape() { return _originalShape; }
 
 bool TadDescriptor::areUnitiesinShape() const { return _unitiesInShape; }
 }  // namespace sd
@@ -80,13 +68,6 @@ size_t hash<sd::TadDescriptor>::operator()(const sd::TadDescriptor &k) const {
 
   // Start with initial hash from unities flag
   uint64_t hash = ModularHasher::hash_scalar(k.areUnitiesinShape());
-
-  // Combine with original shape hash
-  hash = ModularHasher::combine_hashes({
-      hash,
-      std::hash<sd::ShapeDescriptor>()(k.originalShapeConst())
-  });
-
   // Hash the axis vector
   auto& axes = const_cast<sd::TadDescriptor&>(k).axis();
   if (!axes.empty()) {

libnd4j/include/execution/impl/Threads.cpp

@@ -146,6 +146,9 @@ Span2 Span2::build(int loop, uint64_t threadID, uint64_t numThreads, int64_t sta
     default:
       THROW_EXCEPTION("");
   }
+
+  return Span2(startX, stopX, incX, 0, 0, incY);
+
 }
 
 int64_t Span::startX() const {