wehub-resource-sync
124 lines
3.8 KiB
C++
124 lines
3.8 KiB
C++
/*
|
|
* Licensed to the Apache Software Foundation (ASF) under one
|
|
* or more contributor license agreements. See the NOTICE file
|
|
* distributed with this work for additional information
|
|
* regarding copyright ownership. The ASF licenses this file
|
|
* to you under the Apache License, Version 2.0 (the
|
|
* "License"); you may not use this file except in compliance
|
|
* with the License. You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* 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.
|
|
*/
|
|
|
|
/*!
|
|
* \file tvm/ir/scope_stack.h
|
|
* \brief A generic scope stack for managing hierarchical state during IR visiting.
|
|
*/
|
|
#ifndef TVM_IR_SCOPE_STACK_H_
|
|
#define TVM_IR_SCOPE_STACK_H_
|
|
|
|
#include <tvm/ffi/error.h>
|
|
|
|
#include <deque>
|
|
#include <type_traits>
|
|
|
|
namespace tvm {
|
|
|
|
/*!
|
|
* \brief A scope stack for maintaining hierarchical state during IR visiting.
|
|
*
|
|
* During IR tree traversal, visitors often need to track scope-local state
|
|
* (e.g., active constraints, variable bindings) that should be automatically
|
|
* cleaned up when leaving a scope. ScopeStack provides this via WithNewScope,
|
|
* which pushes a new element on entry and pops it on exit.
|
|
*
|
|
* \code
|
|
* ScopeStack<WithGroup<ConstraintContext>> constraints;
|
|
*
|
|
* // In VisitStmt_(ForNode):
|
|
* return constraints.WithNewScope([&]() -> Stmt {
|
|
* constraints.Current().Emplace(analyzer, condition);
|
|
* return StmtExprMutator::VisitStmt_(op);
|
|
* });
|
|
* \endcode
|
|
*
|
|
* \tparam T The element type stored on the stack. Must be default-constructible.
|
|
*/
|
|
template <typename T>
|
|
class ScopeStack {
|
|
public:
|
|
/*! \brief Construct with one initial scope level. */
|
|
ScopeStack() { stack_.emplace_back(); }
|
|
|
|
/*! \brief Return the number of active scopes. */
|
|
size_t size() const { return stack_.size(); }
|
|
|
|
/*! \brief Return true if no scopes are active. */
|
|
bool empty() const { return stack_.empty(); }
|
|
|
|
/*!
|
|
* \brief Access the current (innermost) scope element.
|
|
*
|
|
* The returned reference is stable across push_back/pop_back because
|
|
* std::deque guarantees pointer stability for these operations.
|
|
*
|
|
* \return Mutable reference to the top element.
|
|
*/
|
|
T& Current() {
|
|
TVM_FFI_ICHECK(!stack_.empty());
|
|
return stack_.back();
|
|
}
|
|
|
|
/*! \brief Const access to the current (innermost) scope element. */
|
|
const T& Current() const {
|
|
TVM_FFI_ICHECK(!stack_.empty());
|
|
return stack_.back();
|
|
}
|
|
|
|
/*!
|
|
* \brief Execute body within a new scope.
|
|
*
|
|
* Pushes a new T onto the stack, executes the body, then pops it.
|
|
*
|
|
* \param body A callable to execute within the scope.
|
|
* \return The return value of body(), if non-void.
|
|
*/
|
|
template <typename F>
|
|
auto WithNewScope(F&& body) -> decltype(body()) {
|
|
stack_.emplace_back();
|
|
struct Guard {
|
|
std::deque<T>* stack;
|
|
~Guard() noexcept(false) { stack->pop_back(); }
|
|
} guard{&stack_};
|
|
if constexpr (std::is_void_v<decltype(body())>) {
|
|
body();
|
|
} else {
|
|
return body();
|
|
}
|
|
}
|
|
|
|
private:
|
|
/*!
|
|
* \brief The scope stack.
|
|
*
|
|
* We use std::deque rather than std::vector for pointer stability:
|
|
* references returned by Current() remain valid across push/pop operations.
|
|
* This is critical because methods called on Current() (e.g., Emplace on
|
|
* a WithGroup) may trigger re-entrant code that pushes new scopes onto
|
|
* the same stack. With std::vector the internal buffer reallocation would
|
|
* invalidate the reference, causing use-after-free.
|
|
*/
|
|
std::deque<T> stack_;
|
|
};
|
|
|
|
} // namespace tvm
|
|
|
|
#endif // TVM_IR_SCOPE_STACK_H_
|