chore: import upstream snapshot with attribution

libnd4j/include/memory/AllocationEntry.h

@@ -0,0 +1,52 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// Created by raver119 on 07.05.19.
+//
+
+#ifndef DEV_TESTS_ALLOCATIONENTRY_H
+#define DEV_TESTS_ALLOCATIONENTRY_H
+
+#include <memory/MemoryType.h>
+#include <system/common.h>
+
+#include <string>
+
+namespace sd {
+namespace memory {
+class AllocationEntry {
+ private:
+  MemoryType _memoryType;
+  LongType _pointer;
+  LongType _numBytes;
+  std::string _stack;
+
+ public:
+  AllocationEntry() = default;
+  AllocationEntry(MemoryType type, LongType ptr, LongType numBytes, std::string &stack);
+  ~AllocationEntry() = default;
+
+  LongType numBytes();
+  std::string stackTrace();
+  MemoryType memoryType();
+};
+}  // namespace memory
+}  // namespace sd
+
+#endif  // DEV_TESTS_ALLOCATIONENTRY_H

libnd4j/include/memory/ExternalWorkspace.h

@@ -0,0 +1,52 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+//  @author raver119@gmail.com
+//
+
+#ifndef LIBND4J_EXTERNALWORKSPACE_H
+#define LIBND4J_EXTERNALWORKSPACE_H
+#include <system/common.h>
+
+namespace sd {
+namespace memory {
+class SD_LIB_EXPORT ExternalWorkspace {
+ private:
+  void *_ptrH = nullptr;
+  void *_ptrD = nullptr;
+
+  LongType _sizeH = 0L;
+  LongType _sizeD = 0L;
+
+ public:
+  ExternalWorkspace() = default;
+  ~ExternalWorkspace() = default;
+
+  ExternalWorkspace(Pointer ptrH, LongType sizeH, Pointer ptrD, LongType sizeD);
+
+  void *pointerHost();
+  void *pointerDevice();
+
+  LongType sizeHost();
+  LongType sizeDevice();
+};
+}  // namespace memory
+}  // namespace sd
+
+#endif

libnd4j/include/memory/MemoryCounter.h

@@ -0,0 +1,145 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// @author raver119@gmail.com
+//
+
+#ifndef SD_MEMORYCOUNTER_H
+#define SD_MEMORYCOUNTER_H
+
+#include <memory/MemoryType.h>
+#include <system/common.h>
+
+#include <map>
+#include <mutex>
+
+namespace sd {
+namespace memory {
+/**
+ * This class provides simple per-device counter
+ */
+class SD_LIB_EXPORT MemoryCounter {
+ private:
+  // used for synchronization
+  std::mutex _locker;
+
+  // per-device counters
+  std::map<int, LongType> _deviceCounters;
+
+  // TODO: change this wrt heterogenous stuff on next iteration
+  // per-group counters
+  std::map<MemoryType, LongType> _groupCounters;
+
+  // per-device limits
+  std::map<int, LongType> _deviceLimits;
+
+  // per-group limits
+  std::map<MemoryType, LongType> _groupLimits;
+
+  MemoryCounter();
+  ~MemoryCounter() = default;
+
+ public:
+  static MemoryCounter& getInstance();
+
+  /**
+   * This method checks if allocation of numBytes won't break through per-group or per-device limit
+   * @param numBytes
+   * @return TRUE if allocated ammount will keep us below limit, FALSE otherwise
+   */
+  bool validate(LongType numBytes);
+
+  /**
+   * This method checks if allocation of numBytes won't break through  per-device limit
+   * @param deviceId
+   * @param numBytes
+   * @return TRUE if allocated ammount will keep us below limit, FALSE otherwise
+   */
+  bool validateDevice(int deviceId, LongType numBytes);
+
+  /**
+   * This method checks if allocation of numBytes won't break through per-group limit
+   * @param deviceId
+   * @param numBytes
+   * @return TRUE if allocated ammount will keep us below limit, FALSE otherwise
+   */
+  bool validateGroup(MemoryType group, LongType numBytes);
+
+  /**
+   * This method adds specified number of bytes to specified counter
+   * @param deviceId
+   * @param numBytes
+   */
+  void countIn(int deviceId, LongType numBytes);
+  void countIn(MemoryType group, LongType numBytes);
+
+  /**
+   * This method subtracts specified number of bytes from specified counter
+   * @param deviceId
+   * @param numBytes
+   */
+  void countOut(int deviceId, LongType numBytes);
+  void countOut(MemoryType group, LongType numBytes);
+
+  /**
+   * This method returns amount of memory allocated on specified device
+   * @param deviceId
+   * @return
+   */
+  LongType allocatedDevice(int deviceId);
+
+  /**
+   * This method returns amount of memory allocated in specified group of devices
+   * @param group
+   * @return
+   */
+  LongType allocatedGroup(MemoryType group);
+
+  /**
+   * This method allows to set per-device memory limits
+   * @param deviceId
+   * @param numBytes
+   */
+  void setDeviceLimit(int deviceId, LongType numBytes);
+
+  /**
+   * This method returns current device limit in bytes
+   * @param deviceId
+   * @return
+   */
+  LongType deviceLimit(int deviceId);
+
+  /**
+   * This method allows to set per-group memory limits
+   * @param group
+   * @param numBytes
+   */
+  void setGroupLimit(MemoryType group, LongType numBytes);
+
+  /**
+   * This method returns current group limit in bytes
+   * @param group
+   * @return
+   */
+  LongType groupLimit(MemoryType group);
+};
+}  // namespace memory
+}  // namespace sd
+
+#endif  // SD_MEMORYCOUNTER_H

libnd4j/include/memory/MemoryRegistrator.h

@@ -0,0 +1,72 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// Created by raver119 on 12.09.17.
+//
+
+#ifndef LIBND4J_MEMORYREGISTRATOR_H
+#define LIBND4J_MEMORYREGISTRATOR_H
+
+#include <system/common.h>
+#include <system/op_boilerplate.h>
+
+#include <map>
+#include <mutex>
+#include <unordered_map>
+
+#include "Workspace.h"
+
+namespace sd {
+namespace memory {
+class SD_LIB_EXPORT MemoryRegistrator {
+ protected:
+  Workspace* _workspace;
+  SD_MAP_IMPL<LongType, LongType> _footprint;
+  std::mutex _lock;
+
+  MemoryRegistrator();
+  ~MemoryRegistrator() = default;
+
+ public:
+  static MemoryRegistrator& getInstance();
+  bool hasWorkspaceAttached();
+  Workspace* getWorkspace();
+  void attachWorkspace(Workspace* workspace);
+  void forgetWorkspace();
+
+  /**
+   * This method allows you to set memory requirements for given graph
+   */
+  void setGraphMemoryFootprint(LongType hash, LongType bytes);
+
+  /**
+   * This method allows you to set memory requirements for given graph, ONLY if
+   * new amount of bytes is greater then current one
+   */
+  void setGraphMemoryFootprintIfGreater(LongType hash, LongType bytes);
+
+  /**
+   * This method returns memory requirements for given graph
+   */
+  LongType getGraphMemoryFootprint(LongType hash);
+};
+}  // namespace memory
+}  // namespace sd
+
+#endif  // LIBND4J_MEMORYREGISTRATOR_H

libnd4j/include/memory/MemoryReport.h

@@ -0,0 +1,55 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// Created by raver119 on 11.10.2017.
+//
+
+#ifndef LIBND4J_MEMORYREPORT_H
+#define LIBND4J_MEMORYREPORT_H
+
+#include <system/common.h>
+
+namespace sd {
+namespace memory {
+class SD_LIB_EXPORT MemoryReport {
+ private:
+  LongType _vm = 0;
+  LongType _rss = 0;
+
+ public:
+  MemoryReport() = default;
+  ~MemoryReport() = default;
+
+  bool operator<(const MemoryReport& other) const;
+  bool operator<=(const MemoryReport& other) const;
+  bool operator>(const MemoryReport& other) const;
+  bool operator>=(const MemoryReport& other) const;
+  bool operator==(const MemoryReport& other) const;
+  bool operator!=(const MemoryReport& other) const;
+
+  LongType getVM() const;
+  void setVM(LongType vm);
+
+  LongType getRSS() const;
+  void setRSS(LongType rss);
+};
+}  // namespace memory
+}  // namespace sd
+
+#endif  // LIBND4J_MEMORYREPORT_H

libnd4j/include/memory/MemoryTracker.h

@@ -0,0 +1,59 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// Created by raver119 on 07.05.19.
+//
+
+#ifndef DEV_TESTS_MEMORYTRACKER_H
+#define DEV_TESTS_MEMORYTRACKER_H
+#include <system/common.h>
+
+#include <map>
+#include <mutex>
+#include <string>
+
+#include "AllocationEntry.h"
+
+namespace sd {
+namespace memory {
+/**
+ * This class is used for tracking memory allocation wrt their allocation points in code
+ */
+class SD_LIB_EXPORT MemoryTracker {
+ private:
+  std::map<LongType, AllocationEntry> _allocations;
+  std::map<LongType, AllocationEntry> _released;
+  std::mutex _locker;
+
+  MemoryTracker();
+  ~MemoryTracker() = default;
+
+ public:
+  static MemoryTracker& getInstance();
+
+  void countIn(MemoryType type, Pointer ptr, LongType numBytes);
+  void countOut(Pointer ptr);
+
+  void summarize();
+  void reset();
+};
+}  // namespace memory
+}  // namespace sd
+
+#endif  // DEV_TESTS_MEMORYTRACKER_H

libnd4j/include/memory/MemoryType.h

@@ -0,0 +1,17 @@
+//
+// Created by raver119 on 07.05.19.
+//
+
+#ifndef DEV_TESTS_MEMORYTYPE_H
+#define DEV_TESTS_MEMORYTYPE_H
+
+namespace sd {
+namespace memory {
+enum MemoryType {
+  HOST = 0,
+  DEVICE = 10,
+};
+}
+}  // namespace sd
+
+#endif  // DEV_TESTS_MEMORYTYPE_H

libnd4j/include/memory/MemoryUtils.h

@@ -0,0 +1,37 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// Created by raver119 on 11.10.2017.
+//
+
+#ifndef LIBND4J_MEMORYUTILS_H
+#define LIBND4J_MEMORYUTILS_H
+
+#include "MemoryReport.h"
+
+namespace sd {
+namespace memory {
+class SD_LIB_EXPORT MemoryUtils {
+ public:
+  static bool retrieveMemoryStatistics(MemoryReport& report);
+};
+}  // namespace memory
+}  // namespace sd
+
+#endif  // LIBND4J_MEMORYUTILS_H

libnd4j/include/memory/Workspace.h

@@ -0,0 +1,111 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// This class implements Workspace functionality in c++
+//
+//
+// @author raver119@gmail.com
+//
+
+#ifndef LIBND4J_WORKSPACE_H
+#define LIBND4J_WORKSPACE_H
+
+#include <memory/ExternalWorkspace.h>
+#include <memory/MemoryType.h>
+#include <system/common.h>
+#include <types/float16.h>
+
+#include <atomic>
+#include <mutex>
+#include <vector>
+
+namespace sd {
+namespace memory {
+
+class SD_LIB_EXPORT Workspace {
+ protected:
+  char* _ptrHost = nullptr;
+  char* _ptrDevice = nullptr;
+
+  bool _allocatedHost = false;
+  bool _allocatedDevice = false;
+
+  std::atomic<sd::LongType> _offset;
+  std::atomic<sd::LongType> _offsetSecondary;
+
+  sd::LongType _initialSize = 0L;
+  sd::LongType _initialSizeSecondary = 0L;
+
+  sd::LongType _currentSize = 0L;
+  sd::LongType _currentSizeSecondary = 0L;
+
+  std::mutex _mutexAllocation;
+  std::mutex _mutexSpills;
+
+  bool _externalized = false;
+
+  std::vector<void*> _spills;
+  std::vector<void*> _spillsSecondary;
+
+  std::atomic<sd::LongType> _spillsSize;
+  std::atomic<sd::LongType> _cycleAllocations;
+
+  std::atomic<sd::LongType> _spillsSizeSecondary;
+  std::atomic<sd::LongType> _cycleAllocationsSecondary;
+
+  void init(sd::LongType primaryBytes, sd::LongType secondaryBytes = 0L);
+  void freeSpills();
+
+ public:
+  explicit Workspace(ExternalWorkspace* external);
+  Workspace(sd::LongType initialSize = 0L, sd::LongType secondaryBytes = 0L);
+  ~Workspace();
+
+  sd::LongType getAllocatedSize();
+  sd::LongType getCurrentSize();
+  sd::LongType getCurrentOffset();
+  sd::LongType getSpilledSize();
+  sd::LongType getUsedSize();
+
+  sd::LongType getAllocatedSecondarySize();
+  sd::LongType getCurrentSecondarySize();
+  sd::LongType getCurrentSecondaryOffset();
+  sd::LongType getSpilledSecondarySize();
+  sd::LongType getUsedSecondarySize();
+
+  void expandBy(sd::LongType primaryBytes, sd::LongType secondaryBytes = 0L);
+  void expandTo(sd::LongType primaryBytes, sd::LongType secondaryBytes = 0L);
+
+  //            bool resizeSupported();
+
+  void* allocateBytes(sd::LongType numBytes);
+  void* allocateBytes(MemoryType type, sd::LongType numBytes);
+
+  void scopeIn();
+  void scopeOut();
+
+  /*
+   * This method creates NEW workspace of the same memory size and returns pointer to it
+   */
+  Workspace* clone();
+};
+}  // namespace memory
+}  // namespace sd
+
+#endif  // LIBND4J_WORKSPACE_H

libnd4j/include/memory/cpu/Workspace.cpp

@@ -0,0 +1,192 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// CPU workspaces implementation
+//
+// @author raver119@gmail.com
+//
+
+#include "../Workspace.h"
+
+#include <helpers/logger.h>
+#include <math/templatemath.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <system/op_boilerplate.h>
+
+#include <atomic>
+#include <cstring>
+
+namespace sd {
+namespace memory {
+Workspace::Workspace(ExternalWorkspace *external) {
+  if (external->sizeHost() > 0) {
+    _ptrHost = (char *)external->pointerHost();
+    _ptrDevice = (char *)external->pointerDevice();
+
+    _initialSize = external->sizeHost();
+    _currentSize = external->sizeHost();
+    _offset = 0L;
+    _offsetSecondary = 0L;
+    this->_cycleAllocations = 0;
+    this->_spillsSize = 0;
+
+    _externalized = true;
+  }
+};
+
+Workspace::Workspace(sd::LongType initialSize, sd::LongType secondaryBytes) {
+  if (initialSize > 0) {
+    this->_ptrHost = (char *)malloc(initialSize);
+
+    CHECK_ALLOC(this->_ptrHost, "Failed to allocate new workspace", initialSize);
+
+    memset(this->_ptrHost, 0, initialSize);
+    this->_allocatedHost = true;
+  } else
+    this->_allocatedHost = false;
+
+  this->_initialSize = initialSize;
+  this->_currentSize = initialSize;
+  this->_currentSizeSecondary = 0;
+  this->_spillsSizeSecondary = 0;
+  this->_offset = 0;
+  this->_offsetSecondary = 0;
+  this->_cycleAllocations = 0;
+  this->_spillsSize = 0;
+}
+
+void Workspace::init(sd::LongType bytes, sd::LongType secondaryBytes) {
+  if (this->_currentSize < bytes) {
+    if (this->_allocatedHost && !_externalized) free((void *)this->_ptrHost);
+
+    this->_ptrHost = (char *)malloc(bytes);
+
+    CHECK_ALLOC(this->_ptrHost, "Failed to allocate new workspace", bytes);
+
+    memset(this->_ptrHost, 0, bytes);
+    this->_currentSize = bytes;
+    this->_allocatedHost = true;
+  }
+}
+
+void Workspace::expandBy(sd::LongType numBytes, sd::LongType secondaryBytes) {
+  this->init(_currentSize + numBytes, _currentSizeSecondary + secondaryBytes);
+}
+
+void Workspace::expandTo(sd::LongType numBytes, sd::LongType secondaryBytes) { this->init(numBytes, secondaryBytes); }
+
+void Workspace::freeSpills() {
+  _spillsSize = 0;
+
+  if (_spills.size() < 1) return;
+
+  for (auto v : _spills) free(v);
+
+  _spills.clear();
+}
+
+Workspace::~Workspace() {
+  if (this->_allocatedHost && !_externalized) free((void *)this->_ptrHost);
+
+  freeSpills();
+}
+
+sd::LongType Workspace::getUsedSize() { return getCurrentOffset(); }
+
+sd::LongType Workspace::getCurrentSize() { return _currentSize; }
+
+sd::LongType Workspace::getCurrentOffset() { return _offset.load(); }
+
+void *Workspace::allocateBytes(sd::LongType numBytes) {
+  if (numBytes < 1) THROW_EXCEPTION(allocation_exception::build("Number of bytes for allocation should be positive", numBytes).what());
+
+  // numBytes += 32;
+  void *result = nullptr;
+  this->_cycleAllocations += numBytes;
+  this->_mutexAllocation.lock();
+
+  if (_offset.load() + numBytes > _currentSize) {
+    sd_debug("Allocating %lld bytes in spills\n", numBytes);
+    this->_mutexAllocation.unlock();
+#if defined(SD_ALIGNED_ALLOC)
+    void *p = aligned_alloc(SD_DESIRED_ALIGNMENT, (numBytes + SD_DESIRED_ALIGNMENT - 1) & (-SD_DESIRED_ALIGNMENT));
+#else
+    void *p = malloc(numBytes);
+#endif
+    CHECK_ALLOC(p, "Failed to allocate new workspace", numBytes);
+
+    _mutexSpills.lock();
+    _spills.push_back(p);
+    _mutexSpills.unlock();
+
+    _spillsSize += numBytes;
+
+    return p;
+  }
+
+  result = (void *)(_ptrHost + _offset.load());
+  _offset += numBytes;
+  // memset(result, 0, (int) numBytes);
+
+  sd_debug("Allocating %lld bytes from workspace; Current PTR: %p; Current offset: %lld\n", numBytes, result,
+           _offset.load());
+
+  this->_mutexAllocation.unlock();
+
+  return result;
+}
+
+sd::LongType Workspace::getAllocatedSize() { return getCurrentSize() + getSpilledSize(); }
+
+void Workspace::scopeIn() {
+  freeSpills();
+  init(_cycleAllocations.load());
+  _cycleAllocations = 0;
+}
+
+void Workspace::scopeOut() {
+  _offset = 0;
+  _offsetSecondary = 0;
+}
+
+sd::LongType Workspace::getSpilledSize() { return _spillsSize.load(); }
+
+void *Workspace::allocateBytes(sd::memory::MemoryType type, sd::LongType numBytes) {
+  if (type == DEVICE) THROW_EXCEPTION("CPU backend doesn't have device memory");
+
+  return this->allocateBytes(numBytes);
+}
+
+sd::LongType Workspace::getAllocatedSecondarySize() { return 0L; }
+
+sd::LongType Workspace::getCurrentSecondarySize() { return 0L; }
+
+sd::LongType Workspace::getCurrentSecondaryOffset() { return 0L; }
+
+sd::LongType Workspace::getSpilledSecondarySize() { return 0L; }
+
+sd::LongType Workspace::getUsedSecondarySize() { return 0L; }
+
+Workspace *Workspace::clone() {
+  // for clone we take whatever is higher: current allocated size, or allocated size of current loop
+  return new Workspace(sd::math::sd_max<sd::LongType>(this->getCurrentSize(), this->_cycleAllocations.load()));
+}
+}  // namespace memory
+}  // namespace sd

libnd4j/include/memory/cuda/Workspace.cu

@@ -0,0 +1,257 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// CUDA workspaces implementation
+//
+// @author raver119@gmail.com
+//
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <exceptions/cuda_exception.h>
+#include <helpers/logger.h>
+#include <math/templatemath.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <system/op_boilerplate.h>
+
+#include <atomic>
+#include <cstring>
+
+#include "../Workspace.h"
+
+namespace sd {
+namespace memory {
+Workspace::Workspace(ExternalWorkspace *external) {
+  if (external->sizeHost() > 0) {
+    _ptrHost = (char *)external->pointerHost();
+    _ptrDevice = (char *)external->pointerDevice();
+
+    _initialSize = external->sizeDevice();
+    _currentSize = external->sizeDevice();
+    _initialSizeSecondary = external->sizeHost();
+    _currentSizeSecondary = external->sizeHost();
+    _offset = 0L;
+    _offsetSecondary = 0L;
+    this->_cycleAllocations = 0;
+    this->_cycleAllocationsSecondary = 0;
+    this->_spillsSize = 0;
+    this->_spillsSizeSecondary = 0;
+
+    _externalized = true;
+  }
+}
+
+Workspace::Workspace(LongType primarySize, LongType secondarySize) {
+  if (secondarySize > 0) {
+    auto res = cudaHostAlloc(reinterpret_cast<void **>(&_ptrHost), secondarySize, cudaHostAllocDefault);
+    if (res != 0) throw cuda_exception::build("Can't allocate [HOST] memory", res);
+
+    cudaMemset(this->_ptrHost, 0, secondarySize);
+    this->_allocatedHost = true;
+  } else
+    this->_allocatedHost = false;
+
+  if (primarySize > 0) {
+    auto res = cudaMalloc(reinterpret_cast<void **>(&_ptrDevice), primarySize);
+    if (res != 0) throw cuda_exception::build("Can't allocate [DEVICE] memory", res);
+
+    cudaMemset(this->_ptrDevice, 0, primarySize);
+    this->_allocatedDevice = true;
+  } else
+    this->_allocatedDevice = false;
+
+  this->_initialSize = primarySize;
+  this->_initialSizeSecondary = secondarySize;
+  this->_currentSize = primarySize;
+  this->_currentSizeSecondary = secondarySize;
+  this->_offset = 0;
+  this->_offsetSecondary = 0;
+  this->_cycleAllocations = 0;
+  this->_spillsSize = 0;
+  this->_spillsSizeSecondary = 0;
+}
+
+void Workspace::init(LongType primaryBytes, LongType secondaryBytes) {
+  if (this->_currentSize < primaryBytes) {
+    if (this->_allocatedDevice && !_externalized) cudaFree((void *)this->_ptrDevice);
+
+    auto res = cudaMalloc(reinterpret_cast<void **>(&_ptrDevice), secondaryBytes);
+    if (res != 0) throw cuda_exception::build("Can't allocate [DEVICE] memory", res);
+
+    cudaMemset(this->_ptrDevice, 0, primaryBytes);
+    this->_currentSize = primaryBytes;
+    this->_allocatedDevice = true;
+  }
+
+  if (this->_currentSizeSecondary < secondaryBytes) {
+    if (this->_allocatedHost && !_externalized) cudaFreeHost((void *)this->_ptrHost);
+
+    auto res = cudaHostAlloc(reinterpret_cast<void **>(&_ptrHost), secondaryBytes, cudaHostAllocDefault);
+    if (res != 0) throw cuda_exception::build("Can't allocate [HOST] memory", res);
+
+    cudaMemset(this->_ptrHost, 0, secondaryBytes);
+    this->_currentSizeSecondary = secondaryBytes;
+    this->_allocatedHost = true;
+  }
+}
+
+void Workspace::expandBy(LongType numBytes, LongType secondaryBytes) {
+  this->init(_currentSize + numBytes, _currentSizeSecondary + secondaryBytes);
+}
+
+void Workspace::expandTo(LongType numBytes, LongType secondaryBytes) { this->init(numBytes, secondaryBytes); }
+
+void Workspace::freeSpills() {
+  _spillsSize = 0;
+  _spillsSizeSecondary = 0;
+
+  for (auto v : _spills) cudaFree(v);
+
+  for (auto v : _spillsSecondary) cudaFreeHost(v);
+
+  _spills.clear();
+  _spillsSecondary.clear();
+}
+
+Workspace::~Workspace() {
+  if (this->_allocatedHost && !_externalized) cudaFreeHost((void *)this->_ptrHost);
+
+  if (this->_allocatedDevice && !_externalized) cudaFree((void *)this->_ptrDevice);
+
+  freeSpills();
+}
+
+LongType Workspace::getUsedSize() { return getCurrentOffset(); }
+
+LongType Workspace::getCurrentSize() { return _currentSize; }
+
+LongType Workspace::getCurrentOffset() { return _offset.load(); }
+
+void *Workspace::allocateBytes(LongType numBytes) { return allocateBytes(HOST, numBytes); }
+
+LongType Workspace::getAllocatedSize() { return getCurrentSize() + getSpilledSize(); }
+
+void Workspace::scopeIn() {
+  freeSpills();
+  init(_cycleAllocations.load());
+  _cycleAllocations = 0;
+}
+
+void Workspace::scopeOut() { _offset = 0; }
+
+LongType Workspace::getSpilledSize() { return _spillsSize.load(); }
+
+void *Workspace::allocateBytes(MemoryType type, LongType numBytes) {
+  switch (type) {
+    case HOST: {
+      if (numBytes < 1)
+        throw allocation_exception::build("Number of [HOST] bytes for allocation should be positive", numBytes);
+
+      // numBytes += 32;
+      void *result = nullptr;
+      this->_cycleAllocationsSecondary += numBytes;
+      this->_mutexAllocation.lock();
+
+      if (_offsetSecondary.load() + numBytes > _currentSizeSecondary) {
+        sd_debug("Allocating %lld [HOST] bytes in spills\n", numBytes);
+        this->_mutexAllocation.unlock();
+
+        Pointer p;
+        auto res = cudaHostAlloc(reinterpret_cast<void **>(&p), numBytes, cudaHostAllocDefault);
+        if (res != 0) throw cuda_exception::build("Can't allocate [HOST] memory", res);
+
+        _mutexSpills.lock();
+        _spillsSecondary.push_back(p);
+        _mutexSpills.unlock();
+
+        _spillsSizeSecondary += numBytes;
+
+        return p;
+      }
+
+      result = (void *)(_ptrHost + _offsetSecondary.load());
+      _offsetSecondary += numBytes;
+      // memset(result, 0, (int) numBytes);
+
+      sd_debug("Allocating %lld bytes from [HOST] workspace; Current PTR: %p; Current offset: %lld\n", numBytes, result,
+               _offset.load());
+
+      this->_mutexAllocation.unlock();
+
+      return result;
+    } break;
+    case DEVICE: {
+      if (numBytes < 1)
+        throw allocation_exception::build("Number of [DEVICE] bytes for allocation should be positive", numBytes);
+
+      // numBytes += 32;
+      void *result = nullptr;
+      this->_cycleAllocations += numBytes;
+      this->_mutexAllocation.lock();
+
+      if (_offset.load() + numBytes > _currentSize) {
+        sd_debug("Allocating %lld [DEVICE] bytes in spills\n", numBytes);
+        this->_mutexAllocation.unlock();
+
+        Pointer p;
+        auto res = cudaMalloc(reinterpret_cast<void **>(&p), numBytes);
+        if (res != 0) throw cuda_exception::build("Can't allocate [DEVICE] memory", res);
+
+        _mutexSpills.lock();
+        _spills.push_back(p);
+        _mutexSpills.unlock();
+
+        _spillsSize += numBytes;
+
+        return p;
+      }
+
+      result = (void *)(_ptrDevice + _offset.load());
+      _offset += numBytes;
+      // memset(result, 0, (int) numBytes);
+
+      sd_debug("Allocating %lld bytes from [DEVICE] workspace; Current PTR: %p; Current offset: %lld\n", numBytes,
+               result, _offset.load());
+
+      this->_mutexAllocation.unlock();
+
+      return result;
+    } break;
+    default:
+      THROW_EXCEPTION("Unknown MemoryType was passed in");
+  }
+}
+
+Workspace *Workspace::clone() {
+  // for clone we take whatever is higher: current allocated size, or allocated size of current loop
+  return new Workspace(sd::math::sd_max<LongType>(this->getCurrentSize(), this->_cycleAllocations.load()));
+}
+
+LongType Workspace::getAllocatedSecondarySize() { return getCurrentSecondarySize() + getSpilledSecondarySize(); }
+
+LongType Workspace::getCurrentSecondarySize() { return _currentSizeSecondary; }
+
+LongType Workspace::getCurrentSecondaryOffset() { return _offsetSecondary.load(); }
+
+LongType Workspace::getSpilledSecondarySize() { return _spillsSizeSecondary; }
+
+LongType Workspace::getUsedSecondarySize() { return getCurrentSecondaryOffset(); }
+
+}  // namespace memory
+}  // namespace sd

libnd4j/include/memory/impl/AllocationEntry.cpp

@@ -0,0 +1,39 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// Created by raver119 on 07.05.19.
+//
+#include <memory/AllocationEntry.h>
+
+namespace sd {
+namespace memory {
+AllocationEntry::AllocationEntry(MemoryType type, LongType ptr, LongType numBytes, std::string &stack) {
+  _pointer = ptr;
+  _numBytes = numBytes;
+  _stack = stack;
+  _memoryType = type;
+}
+
+std::string AllocationEntry::stackTrace() { return _stack; }
+
+LongType AllocationEntry::numBytes() { return _numBytes; }
+
+MemoryType AllocationEntry::memoryType() { return _memoryType; }
+}  // namespace memory
+}  // namespace sd

libnd4j/include/memory/impl/ExternalWorkspace.cpp

@@ -0,0 +1,42 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+//  @author raver119@gmail.com
+//
+#include <memory/ExternalWorkspace.h>
+
+namespace sd {
+namespace memory {
+ExternalWorkspace::ExternalWorkspace(Pointer ptrH, LongType sizeH, Pointer ptrD, LongType sizeD) {
+  _ptrH = ptrH;
+  _sizeH = sizeH;
+
+  _ptrD = ptrD;
+  _sizeD = sizeD;
+};
+
+void* ExternalWorkspace::pointerHost() { return _ptrH; }
+
+void* ExternalWorkspace::pointerDevice() { return _ptrD; }
+
+LongType ExternalWorkspace::sizeHost() { return _sizeH; }
+
+LongType ExternalWorkspace::sizeDevice() { return _sizeD; }
+}  // namespace memory
+}  // namespace sd

libnd4j/include/memory/impl/MemoryCounter.cpp

@@ -0,0 +1,129 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// @author raver119@gmail.com
+//
+#include "../MemoryCounter.h"
+
+#include <execution/AffinityManager.h>
+#include <helpers/logger.h>
+#include <system/Environment.h>
+
+namespace sd {
+namespace memory {
+
+MemoryCounter::MemoryCounter() {
+  auto numDevices = AffinityManager::numberOfDevices();
+
+  // setting default 0s
+  for (int e = 0; e < numDevices; e++) {
+    _deviceLimits[e] = 0;
+    _deviceCounters[e] = 0;
+  }
+
+  // setting initial values for limits
+  _groupLimits[HOST] = Environment::getInstance().maxPrimaryMemory();
+  _groupLimits[DEVICE] = Environment::getInstance().maxSpecialMemory();
+
+  // setting initial counter values
+  _groupCounters[HOST] = 0;
+  _groupCounters[DEVICE] = 0;
+}
+
+MemoryCounter& MemoryCounter::getInstance() {
+  static MemoryCounter instance;
+  return instance;
+}
+
+void MemoryCounter::countIn(int deviceId, LongType numBytes) {
+  std::lock_guard<std::mutex> lock(_locker);
+  _deviceCounters[deviceId] += numBytes;
+}
+
+void MemoryCounter::countIn(MemoryType group, LongType numBytes) {
+  std::lock_guard<std::mutex> lock(_locker);
+  _groupCounters[group] += numBytes;
+}
+
+void MemoryCounter::countOut(int deviceId, LongType numBytes) {
+  std::lock_guard<std::mutex> lock(_locker);
+  _deviceCounters[deviceId] -= numBytes;
+}
+
+void MemoryCounter::countOut(MemoryType group, LongType numBytes) {
+  std::lock_guard<std::mutex> lock(_locker);
+  _groupCounters[group] -= numBytes;
+}
+
+bool MemoryCounter::validate(LongType numBytes) {
+  auto deviceId = AffinityManager::currentDeviceId();
+  return validateDevice(deviceId, numBytes);
+}
+
+bool MemoryCounter::validateDevice(int deviceId, LongType numBytes) {
+  std::lock_guard<std::mutex> lock(_locker);
+  auto dLimit = _deviceLimits[deviceId];
+  if (dLimit <= 0) return true;
+
+  auto dAlloc = _deviceCounters[deviceId];
+
+  return numBytes + dAlloc <= dLimit;
+}
+
+bool MemoryCounter::validateGroup(MemoryType group, LongType numBytes) {
+  std::lock_guard<std::mutex> lock(_locker);
+  auto gLimit = _groupLimits[group];
+  if (gLimit <= 0) return true;
+
+  auto gAlloc = _groupCounters[group];
+
+  return numBytes + gAlloc <= gLimit;
+}
+
+LongType MemoryCounter::allocatedDevice(int deviceId) {
+  std::lock_guard<std::mutex> lock(_locker);
+  return _deviceCounters[deviceId];
+}
+
+LongType MemoryCounter::allocatedGroup(MemoryType group) {
+  std::lock_guard<std::mutex> lock(_locker);
+  return _groupCounters[group];
+}
+
+void MemoryCounter::setDeviceLimit(int deviceId, LongType numBytes) {
+  std::lock_guard<std::mutex> lock(_locker);
+  _deviceLimits[deviceId] = numBytes;
+}
+
+void MemoryCounter::setGroupLimit(MemoryType group, LongType numBytes) {
+  std::lock_guard<std::mutex> lock(_locker);
+  _groupLimits[group] = numBytes;
+}
+
+LongType MemoryCounter::deviceLimit(int deviceId) {
+  std::lock_guard<std::mutex> lock(_locker);
+  return _deviceLimits[deviceId];
+}
+
+LongType MemoryCounter::groupLimit(MemoryType group) {
+  std::lock_guard<std::mutex> lock(_locker);
+  return _groupLimits[group];
+}
+}  // namespace memory
+}  // namespace sd

libnd4j/include/memory/impl/MemoryRegistrator.cpp

@@ -0,0 +1,74 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// Created by raver119 on 07.10.2017.
+//
+#include <memory/MemoryRegistrator.h>
+
+namespace sd {
+namespace memory {
+
+MemoryRegistrator::MemoryRegistrator() { _workspace = nullptr; };
+
+MemoryRegistrator& MemoryRegistrator::getInstance() {
+  static MemoryRegistrator instance;
+  return instance;
+}
+
+bool MemoryRegistrator::hasWorkspaceAttached() { return _workspace != nullptr; }
+
+Workspace* MemoryRegistrator::getWorkspace() { return _workspace; }
+
+void MemoryRegistrator::attachWorkspace(Workspace* workspace) { _workspace = workspace; }
+
+void MemoryRegistrator::forgetWorkspace() { _workspace = nullptr; }
+
+void MemoryRegistrator::setGraphMemoryFootprint(LongType hash, LongType bytes) {
+  _lock.lock();
+
+  _footprint[hash] = bytes;
+
+  _lock.unlock();
+}
+
+void MemoryRegistrator::setGraphMemoryFootprintIfGreater(LongType hash, LongType bytes) {
+  _lock.lock();
+
+  if (_footprint.count(hash) == 0)
+    _footprint[hash] = bytes;
+  else {
+    LongType cv = _footprint[hash];
+    if (bytes > cv) _footprint[hash] = bytes;
+  }
+
+  _lock.unlock();
+}
+
+LongType MemoryRegistrator::getGraphMemoryFootprint(LongType hash) {
+  _lock.lock();
+
+  LongType result = 0L;
+  if (_footprint.count(hash) > 0) result = _footprint[hash];
+
+  _lock.unlock();
+
+  return result;
+}
+}  // namespace memory
+}  // namespace sd

libnd4j/include/memory/impl/MemoryReport.cpp

@@ -0,0 +1,54 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// Created by raver119 on 11.10.2017.
+//
+#include "memory/MemoryReport.h"
+
+bool sd::memory::MemoryReport::operator<(const MemoryReport &other) const {
+  return this->_rss < other._rss;
+}
+
+bool sd::memory::MemoryReport::operator>(const MemoryReport &other) const {
+  return this->_rss > other._rss;
+}
+
+bool sd::memory::MemoryReport::operator==(const MemoryReport &other) const {
+  return this->_rss == other._rss;
+}
+
+bool sd::memory::MemoryReport::operator!=(const MemoryReport &other) const {
+  return this->_rss != other._rss;
+}
+
+bool sd::memory::MemoryReport::operator<=(const MemoryReport &other) const {
+  return this->_rss <= other._rss;
+}
+
+bool sd::memory::MemoryReport::operator>=(const MemoryReport &other) const {
+  return this->_rss >= other._rss;
+}
+
+sd::LongType sd::memory::MemoryReport::getVM() const { return _vm; }
+
+void sd::memory::MemoryReport::setVM(LongType _vm) { MemoryReport::_vm = _vm; }
+
+sd::LongType sd::memory::MemoryReport::getRSS() const { return _rss; }
+
+void sd::memory::MemoryReport::setRSS(LongType _rss) { MemoryReport::_rss = _rss; }

libnd4j/include/memory/impl/MemoryTracker.cpp

@@ -0,0 +1,169 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// Created by raver119 on 07.05.19.
+//
+#include <helpers/logger.h>
+#include <memory/MemoryTracker.h>
+#include <stdlib.h>
+
+#include <stdexcept>
+
+#if defined(__GNUC__) && !defined(SD_WINDOWS_BUILD) && !defined(__MINGW64__) && !defined(SD_ANDROID_BUILD) && \
+    !defined(SD_IOS_BUILD) && !defined(SD_APPLE_BUILD)
+
+#include <cxxabi.h>
+#include <execinfo.h>
+#include <unistd.h>
+
+#endif
+
+namespace sd {
+namespace memory {
+
+MemoryTracker::MemoryTracker() {
+  //
+}
+
+MemoryTracker &MemoryTracker::getInstance() {
+  static MemoryTracker instance;
+  return instance;
+}
+
+#if defined(__GNUC__) && !defined(__CYGWIN__) && !defined(SD_WINDOWS) && !defined(__MINGW64__) && \
+    !defined(SD_ANDROID_BUILD) && !defined(SD_IOS_BUILD) && !defined(SD_APPLE_BUILD)
+std::string demangle(char *message) {
+  char *mangled_name = 0, *offset_begin = 0, *offset_end = 0;
+
+  // find parantheses and +address offset surrounding mangled name
+  for (char *p = message; *p; ++p) {
+    if (*p == '(') {
+      mangled_name = p;
+    } else if (*p == '+') {
+      offset_begin = p;
+    } else if (*p == ')') {
+      offset_end = p;
+      break;
+    }
+  }
+
+  // if the line could be processed, attempt to demangle the symbol
+  if (mangled_name && offset_begin && offset_end && mangled_name < offset_begin) {
+    *mangled_name++ = '\0';
+    *offset_begin++ = '\0';
+    *offset_end++ = '\0';
+
+    int status;
+    char *real_name = abi::__cxa_demangle(mangled_name, 0, 0, &status);
+
+    // if demangling is successful, output the demangled function name
+    if (status == 0) {
+      std::string result(real_name);
+      free(real_name);
+      return result;
+    } else {
+      // otherwise, output the mangled function name
+      std::string result(message);
+      free(real_name);
+      return result;
+    }
+  }
+
+  // safe return
+  return std::string("");
+}
+
+#endif
+
+void MemoryTracker::countIn(MemoryType type, Pointer ptr, LongType numBytes) {
+#if defined(__GNUC__) && !defined(__MINGW64__) && !defined(__CYGWIN__) && !defined(SD_ANDROID_BUILD) && \
+    !defined(SD_WINDOWS) && !defined(SD_IOS_BUILD) && !defined(SD_APPLE_BUILD)
+  if (Environment::getInstance().isDetectingLeaks()) {
+    auto lptr = reinterpret_cast<LongType>(ptr);
+
+    _locker.lock();
+
+    void *array[50];
+    size_t size;
+    char **messages;
+    size = backtrace(array, 50);
+
+    std::string stack("");
+    messages = backtrace_symbols(array, size);
+    for (size_t i = 1; i < size && messages != NULL; ++i) {
+      stack += demangle(messages[i]) + "\n";
+    }
+
+    free(messages);
+
+    if (stack.find("ConstantTad") != std::string::npos || stack.find("ConstantShape") != std::string::npos) {
+      _locker.unlock();
+      return;
+    }
+
+    std::pair<LongType, AllocationEntry> pair(lptr, AllocationEntry(type, lptr, numBytes, stack));
+    _allocations.insert(pair);
+
+    _locker.unlock();
+  }
+#endif
+}
+
+void MemoryTracker::countOut(Pointer ptr) {
+#if defined(__GNUC__) && !defined(__MINGW64__) && !defined(SD_ANDROID_BUILD) && !defined(SD_IOS_BUILD) && \
+    !defined(SD_APPLE_BUILD)
+  if (Environment::getInstance().isDetectingLeaks()) {
+    auto lptr = reinterpret_cast<LongType>(ptr);
+
+    _locker.lock();
+    if (_released.count(lptr) > 0) {
+       THROW_EXCEPTION("Double free!");
+    }
+
+    if (_allocations.count(lptr) > 0) {
+       auto entry = _allocations[lptr];
+       std::pair<sd::LongType, AllocationEntry> pair(lptr, entry);
+      _released.insert(pair);
+      _allocations.erase(lptr);
+    }
+
+    _locker.unlock();
+  }
+#endif
+}
+
+void MemoryTracker::summarize() {
+  if (!_allocations.empty()) {
+    sd_printf("\n%i leaked allocations\n", (int)_allocations.size());
+
+    for (auto &v : _allocations) {
+      sd_printf("Leak of %i [%s] bytes\n%s\n\n", (int)v.second.numBytes(),
+                v.second.memoryType() == MemoryType::HOST ? "HOST" : "DEVICE", v.second.stackTrace().c_str());
+    }
+
+    THROW_EXCEPTION("Non-released allocations found");
+  }
+}
+
+void MemoryTracker::reset() {
+  _allocations.clear();
+  _released.clear();
+}
+}  // namespace memory
+}  // namespace sd

libnd4j/include/memory/impl/MemoryUtils.cpp

@@ -0,0 +1,87 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// Created by raver119 on 11.10.2017.
+//
+#include "memory/MemoryUtils.h"
+
+#include <helpers/logger.h>
+
+#if defined(__APPLE__)
+#include <mach/mach.h>
+#include <sys/resource.h>
+#elif defined(_WIN32) || defined(__WIN32__) || defined(WIN32)
+
+#else
+// linux
+#include <fcntl.h>
+#include <sys/resource.h>
+#include <unistd.h>
+
+#include <cstring>
+#endif
+
+bool sd::memory::MemoryUtils::retrieveMemoryStatistics(MemoryReport &report) {
+#if defined(__APPLE__)
+  sd_debug("APPLE route\n", "");
+  /*
+      struct task_basic_info t_info;
+      mach_msg_type_number_t t_info_count = TASK_BASIC_INFO_COUNT;
+
+      if (KERN_SUCCESS != task_info(mach_task_self(), TASK_BASIC_INFO, (task_info_t)&t_info, &t_info_count))
+          return false;
+
+      report.setVM(t_info.resident_size);
+      report.setRSS(t_info.resident_size);
+
+
+      sd_debug("RSS: %lld; VM: %lld;\n", report.getRSS(), report.getVM());
+  */
+  struct rusage _usage;
+
+  auto res = getrusage(RUSAGE_SELF, &_usage);
+
+  report.setRSS(_usage.ru_maxrss);
+
+  sd_debug("Usage: %lld; %lld; %lld; %lld;\n", _usage.ru_ixrss, _usage.ru_idrss, _usage.ru_isrss, _usage.ru_maxrss);
+
+  return true;
+#elif defined(_WIN32) || defined(__WIN32__) || defined(WIN32)
+  sd_debug("WIN32 route\n", "");
+
+#else
+  sd_debug("LINUX route\n", "");
+  int fd = open("/proc/self/statm", O_RDONLY, 0);
+  if (fd >= 0) {
+    char line[256];
+    char* s;
+    int n;
+    lseek(fd, 0, SEEK_SET);
+    if ((n = read(fd, line, sizeof(line))) > 0 && (s = (char*)memchr(line, ' ', n)) != NULL) {
+      report.setRSS((LongType)(atoll(s + 1) * getpagesize()));
+    }
+    close(fd);
+  }
+
+
+  return true;
+#endif
+
+  return false;
+}