ai.djl.ndarray.NDScope Maven / Gradle / Ivy
/*
* Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
* with the License. A copy of the License is located at
*
* http://aws.amazon.com/apache2.0/
*
* or in the "license" file accompanying this file. This file 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.
*/
package ai.djl.ndarray;
import java.util.ArrayDeque;
import java.util.Deque;
import java.util.IdentityHashMap;
/**
* A class that tracks {@link NDResource} objects created in the try-with-resource block and close
* them automatically when out of the block scope.
*
* This class has been derived from {@code org.bytedeco.javacpp.PointerScope} by Samuel Audet
*/
public class NDScope implements AutoCloseable {
private static final ThreadLocal> SCOPE_STACK =
ThreadLocal.withInitial(ArrayDeque::new);
private IdentityHashMap resources;
/** Constructs a new {@code NDScope} instance. */
@SuppressWarnings("this-escape")
public NDScope() {
resources = new IdentityHashMap<>();
SCOPE_STACK.get().addLast(this);
}
/**
* Registers {@link NDArray} object to this scope.
*
* @param array the {@link NDArray} object
*/
public static void register(NDArray array) {
Deque queue = SCOPE_STACK.get();
if (queue.isEmpty()) {
return;
}
queue.getLast().resources.put(array, array);
}
/**
* Unregisters {@link NDArray} object from this scope.
*
* @param array the {@link NDArray} object
*/
public static void unregister(NDArray array) {
Deque queue = SCOPE_STACK.get();
if (queue.isEmpty()) {
return;
}
queue.getLast().resources.remove(array);
}
/**
* Unregisters {@link NDArray} object from this scope.
*
* @param arrays the array of {@link NDArray} object
*/
public static void unregister(NDArray... arrays) {
for (NDArray array : arrays) {
unregister(array);
}
}
/**
* Unregisters {@link NDArray} object from this scope.
*
* @param ndlist the {@link NDList} object
*/
public static void unregister(NDList ndlist) {
ndlist.forEach(NDScope::unregister);
}
/** {@inheritDoc} */
@Override
public void close() {
for (NDArray array : resources.keySet()) {
array.close();
}
SCOPE_STACK.get().remove(this);
}
/**
* A method that does nothing.
*
* You may use it if you do not have a better way to suppress the warning of a created but
* not explicitly used scope.
*/
public void suppressNotUsedWarning() {
// do nothing
}
}