Cuda refactoring: (deeplearning4j#10195)

Integrates device based shape buffers for direct shape tries now
Update cmakelists to ignore helpers/cpu when compiling cuda
Remove aurora checks from nd4j backend

Author: agibsonccc

libnd4j/CMakeLists.txt.in

@@ -2,8 +2,6 @@ cmake_minimum_required(VERSION 2.8.2)
 
 project(flatbuffers-download NONE)
 
-# Force clean everything first
-file(REMOVE_RECURSE ${CMAKE_CURRENT_BINARY_DIR}/*)
 
 include(ExternalProject)
 ExternalProject_Add(flatbuffers
@@ -16,7 +14,8 @@ ExternalProject_Add(flatbuffers
         -DCMAKE_C_COMPILER=${CMAKE_C_COMPILER}
         -DFLATBUFFERS_BUILD_FLATC=ON
         -DCMAKE_BUILD_TYPE=Release
-    DOWNLOAD_COMMAND  ${CMAKE_COMMAND} -E rm -rf ${CMAKE_CURRENT_BINARY_DIR}/flatbuffers-src && git clone https://github.com/google/flatbuffers/ ${CMAKE_CURRENT_BINARY_DIR}/flatbuffers-src
+   DOWNLOAD_COMMAND  ${CMAKE_COMMAND} -E rm -rf ${CMAKE_CURRENT_BINARY_DIR}/flatbuffers-src COMMAND git clone https://github.com/google/flatbuffers/ ${CMAKE_CURRENT_BINARY_DIR}/flatbuffers-src
+
     UPDATE_COMMAND    ""
     INSTALL_COMMAND   ""
     TEST_COMMAND      ""

libnd4j/blas/CMakeLists.txt

@@ -356,7 +356,16 @@ if(SD_CUDA)
         file(GLOB_RECURSE CUSTOMOPS_SOURCES ../include/ops/declarable/generic/*.cpp)
         file(GLOB_RECURSE CUSTOMOPS_HELPERS_SOURCES ../include/ops/declarable/helpers/cuda/*.cu ../include/ops/declarable/helpers/impl/*.cpp)
         file(GLOB_RECURSE OPS_SOURCES ../include/ops/impl/*.cpp ../include/ops/declarable/impl/*.cpp ../include/ops/*.h)
-        file(GLOB_RECURSE HELPERS_SOURCES ../include/build_info.cpp ../include/ConstMessages.cpp ../include/helpers/*.cpp ../include/helpers/cuda/*.cu ../include/helpers/*.h)
+        file(GLOB_RECURSE HELPERS_SOURCES
+                ../include/build_info.cpp
+                ../include/ConstMessages.cpp
+                ../include/helpers/*.cpp
+                ../include/helpers/cuda/*.cu
+                ../include/helpers/*.h)
+        file(GLOB CPU_HELPERS_TO_EXCLUDE
+                ../include/helpers/cpu/*.cpp)
+        # remove helpers/cpu
+        list(REMOVE_ITEM HELPERS_SOURCES ${CPU_HELPERS_TO_EXCLUDE})
         file(GLOB_RECURSE INDEXING_SOURCES ../include/indexing/*.cpp ../include/indexing/*.h)
         file(GLOB_RECURSE LOOPS_SOURCES ../include/loops/impl/*.cpp ../include/loops/*.h)
         file(GLOB_RECURSE LEGACY_SOURCES ../include/legacy/impl/*.cpp ../include/legacy/*.cu ../include/legacy/*.h)

libnd4j/include/helpers/DirectShapeTrie.h

@@ -3,7 +3,7 @@
 
 #include <array/ConstantShapeBuffer.h>
 #include <system/common.h>
-
+#include "ShapeBufferPlatformHelper.h"
 #include <array>
 #include <atomic>
 #include <memory>
@@ -78,6 +78,8 @@ class SD_LIB_EXPORT DirectShapeTrie {
       // Make sure mutexes are properly initialized
       new (&_mutexes[i]) SHAPE_MUTEX_TYPE();  // Explicit initialization
     }
+
+    ShapeBufferPlatformHelper::initialize();
   }
 
   // Delete copy constructor and assignment
@@ -88,9 +90,6 @@ class SD_LIB_EXPORT DirectShapeTrie {
   DirectShapeTrie(DirectShapeTrie&&) = delete;
   DirectShapeTrie& operator=(DirectShapeTrie&&) = delete;
 
-  // Create a shape buffer from shapeInfo
-  ConstantShapeBuffer* createBuffer(const LongType* shapeInfo);
-
   // Improved thread-safe getOrCreate
   ConstantShapeBuffer* getOrCreate(const LongType* shapeInfo);
 

libnd4j/include/helpers/ShapeBufferCreator.h

@@ -0,0 +1,48 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * This program and the accompanying materials are made available under the
+ *  * terms of the Apache License, Version 2.0 which is available at
+ *  * https://www.apache.org/licenses/LICENSE-2.0.
+ *  *
+ *  * See the NOTICE file distributed with this work for additional
+ *  * information regarding copyright ownership.
+ *  * Unless required by applicable law or agreed to in writing, software
+ *  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *  * License for the specific language governing permissions and limitations
+ *  * under the License.
+ *  *
+ *  * SPDX-License-Identifier: Apache-2.0
+ *  *****************************************************************************
+ */
+
+#ifndef LIBND4J_SHAPEBUFFERCREATOR_H
+#define LIBND4J_SHAPEBUFFERCREATOR_H
+
+#include "array/ConstantShapeBuffer.h"
+
+namespace sd {
+
+/**
+ * Interface for creating ConstantShapeBuffer objects.
+ * This allows for platform-specific implementations (CPU, CUDA, etc).
+ */
+class ShapeBufferCreator {
+public:
+    virtual ~ShapeBufferCreator() = default;
+
+    /**
+     * Create a ConstantShapeBuffer from the given shape information
+     * 
+     * @param shapeInfo Pointer to shape information
+     * @param rank Rank of the shape
+     * @return A new ConstantShapeBuffer instance
+     */
+    virtual ConstantShapeBuffer* create(const LongType* shapeInfo, int rank) = 0;
+};
+
+} // namespace sd
+
+#endif // LIBND4J_SHAPEBUFFERCREATOR_H

libnd4j/include/helpers/ShapeBufferCreatorHelper.h

@@ -0,0 +1,50 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * This program and the accompanying materials are made available under the
+ *  * terms of the Apache License, Version 2.0 which is available at
+ *  * https://www.apache.org/licenses/LICENSE-2.0.
+ *  *
+ *  * See the NOTICE file distributed with this work for additional
+ *  * information regarding copyright ownership.
+ *  * Unless required by applicable law or agreed to in writing, software
+ *  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *  * License for the specific language governing permissions and limitations
+ *  * under the License.
+ *  *
+ *  * SPDX-License-Identifier: Apache-2.0
+ *  *****************************************************************************
+ */
+
+#ifndef LIBND4J_SHAPEBUFFERCREATORHELPER_H
+#define LIBND4J_SHAPEBUFFERCREATORHELPER_H
+
+#include <helpers/ShapeBufferCreator.h>
+#include <exception>
+
+namespace sd {
+
+/**
+ * Helper class to manage ShapeBufferCreator instances and provide global access
+ */
+class ShapeBufferCreatorHelper {
+public:
+    /**
+     * Get the current ShapeBufferCreator instance
+     */
+    static ShapeBufferCreator& getCurrentCreator();
+    
+    /**
+     * Set the current ShapeBufferCreator to use
+     */
+    static void setCurrentCreator(ShapeBufferCreator* creator);
+    
+private:
+    static ShapeBufferCreator* currentCreator_;
+};
+
+} // namespace sd
+
+#endif // LIBND4J_SHAPEBUFFERCREATORHELPER_H

libnd4j/include/helpers/ShapeBufferPlatformHelper.h

@@ -0,0 +1,53 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * This program and the accompanying materials are made available under the
+ *  * terms of the Apache License, Version 2.0 which is available at
+ *  * https://www.apache.org/licenses/LICENSE-2.0.
+ *  *
+ *  * See the NOTICE file distributed with this work for additional
+ *  * information regarding copyright ownership.
+ *  * Unless required by applicable law or agreed to in writing, software
+ *  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *  * License for the specific language governing permissions and limitations
+ *  * under the License.
+ *  *
+ *  * SPDX-License-Identifier: Apache-2.0
+ *  *****************************************************************************
+ */
+
+#ifndef LIBND4J_SHAPEBUFFERPLATFORMHELPER_H
+#define LIBND4J_SHAPEBUFFERPLATFORMHELPER_H
+
+#include <helpers/ShapeBufferCreatorHelper.h>
+
+namespace sd {
+
+/**
+ * Platform-specific initialization helper
+ * Takes care of setting up the correct creators based on the available hardware
+ */
+class ShapeBufferPlatformHelper {
+public:
+    /**
+     * Initialize platform-specific components
+     * This method should be called during the application startup
+     * to ensure proper creators are set based on the available hardware
+     */
+    static void initialize();
+
+    /**
+     * Automatic initialization through static initialization
+     * C++17 guarantees thread-safety for static initialization
+     */
+    static inline const bool initialized = (initialize(), true);
+
+private:
+ ShapeBufferPlatformHelper() = delete;  // Prevent instantiation
+};
+
+} // namespace sd
+
+#endif  // LIBND4J_SHAPEBUFFERPLATFORMHELPER_H

libnd4j/include/helpers/cpu/CpuShapeBufferCreator.cpp

@@ -0,0 +1,47 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * This program and the accompanying materials are made available under the
+ *  * terms of the Apache License, Version 2.0 which is available at
+ *  * https://www.apache.org/licenses/LICENSE-2.0.
+ *  *
+ *  * See the NOTICE file distributed with this work for additional
+ *  * information regarding copyright ownership.
+ *  * Unless required by applicable law or agreed to in writing, software
+ *  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *  * License for the specific language governing permissions and limitations
+ *  * under the License.
+ *  *
+ *  * SPDX-License-Identifier: Apache-2.0
+ *  *****************************************************************************
+ */
+
+#include <helpers/cpu/CpuShapeBufferCreator.h>
+
+
+namespace sd {
+
+ConstantShapeBuffer* CpuShapeBufferCreator::create(const LongType* shapeInfo, int rank) {
+    const int shapeInfoLength = shape::shapeInfoLength(rank);
+    LongType* shapeCopy = new LongType[shapeInfoLength];
+    std::memcpy(shapeCopy, shapeInfo, shapeInfoLength * sizeof(LongType));
+    
+    auto deallocator = std::shared_ptr<PrimaryPointerDeallocator>(
+        new PrimaryPointerDeallocator(),
+        [] (PrimaryPointerDeallocator* ptr) { delete ptr; }
+    );
+    
+    auto hPtr = std::make_shared<PointerWrapper>(shapeCopy, deallocator);
+    auto buffer = new ConstantShapeBuffer(hPtr);
+    
+    return buffer;
+}
+
+CpuShapeBufferCreator& CpuShapeBufferCreator::getInstance() {
+    static CpuShapeBufferCreator instance;
+    return instance;
+}
+
+} // namespace sd

libnd4j/include/helpers/cpu/CpuShapeBufferCreator.h

@@ -0,0 +1,49 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * This program and the accompanying materials are made available under the
+ *  * terms of the Apache License, Version 2.0 which is available at
+ *  * https://www.apache.org/licenses/LICENSE-2.0.
+ *  *
+ *  * See the NOTICE file distributed with this work for additional
+ *  * information regarding copyright ownership.
+ *  * Unless required by applicable law or agreed to in writing, software
+ *  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *  * License for the specific language governing permissions and limitations
+ *  * under the License.
+ *  *
+ *  * SPDX-License-Identifier: Apache-2.0
+ *  *****************************************************************************
+ */
+
+#ifndef LIBND4J_CPUSHAPEBUFFERCREATOR_H
+#define LIBND4J_CPUSHAPEBUFFERCREATOR_H
+
+#include <helpers/ShapeBufferCreator.h>
+#include <helpers/shape.h>
+#include <memory>
+
+namespace sd {
+
+/**
+ * CPU implementation of the ShapeBufferCreator.
+ */
+class CpuShapeBufferCreator : public ShapeBufferCreator {
+public:
+    /**
+     * Create a ConstantShapeBuffer for CPU usage
+     */
+    ConstantShapeBuffer* create(const LongType* shapeInfo, int rank) override;
+    
+    // Singleton pattern implementation
+    static CpuShapeBufferCreator& getInstance();
+
+private:
+    CpuShapeBufferCreator() = default;
+};
+
+} // namespace sd
+
+#endif // LIBND4J_CPUSHAPEBUFFERCREATOR_H

libnd4j/include/helpers/cuda/CudaShapeBufferCreator.cu

@@ -0,0 +1,57 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * This program and the accompanying materials are made available under the
+ *  * terms of the Apache License, Version 2.0 which is available at
+ *  * https://www.apache.org/licenses/LICENSE-2.0.
+ *  *
+ *  * See the NOTICE file distributed with this work for additional
+ *  * information regarding copyright ownership.
+ *  * Unless required by applicable law or agreed to in writing, software
+ *  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *  * License for the specific language governing permissions and limitations
+ *  * under the License.
+ *  *
+ *  * SPDX-License-Identifier: Apache-2.0
+ *  *****************************************************************************
+ */
+
+#include <helpers/ConstantHelper.h>
+#include <helpers/cuda/CudaShapeBufferCreator.h>
+
+#include "array/CudaPointerDeallocator.h"
+#include "array/PrimaryPointerDeallocator.h"
+
+namespace sd {
+
+ConstantShapeBuffer* CudaShapeBufferCreator::create(const LongType* shapeInfo, int rank) {
+    const int shapeInfoLength = shape::shapeInfoLength(rank);
+    LongType* shapeCopy = new LongType[shapeInfoLength];
+    std::memcpy(shapeCopy, shapeInfo, shapeInfoLength * sizeof(LongType));
+    
+    auto deallocator = std::shared_ptr<PrimaryPointerDeallocator>(
+        new PrimaryPointerDeallocator(),
+        [] (PrimaryPointerDeallocator* ptr) { delete ptr; }
+    );
+    
+    auto hPtr = std::make_shared<PointerWrapper>(shapeCopy, deallocator);
+    
+    // Create device pointer for CUDA
+    auto dPtr = std::make_shared<PointerWrapper>(
+        ConstantHelper::getInstance().replicatePointer(hPtr->pointer(),
+                                                   shape::shapeInfoByteLength(hPtr->pointerAsT<sd::LongType>())),
+        std::make_shared<CudaPointerDeallocator>());
+        
+    ConstantShapeBuffer *buffer = new ConstantShapeBuffer(hPtr, dPtr);
+    
+    return buffer;
+}
+
+CudaShapeBufferCreator& CudaShapeBufferCreator::getInstance() {
+    static CudaShapeBufferCreator instance;
+    return instance;
+}
+
+} // namespace sd