Overload get() function for Optional type. (apache#9748)

* upd * simplify * upd * fix * upd * fix docstring
cjia4 · Feb 16, 2022 · 1da0093 · 1da0093
1 parent 75122db
commit 1da0093
Showing 3 changed files with 7 additions and 2 deletions.
diff --git a/include/tvm/runtime/container/optional.h b/include/tvm/runtime/container/optional.h
@@ -93,6 +93,11 @@ class Optional : public ObjectRef {
     ICHECK(data_ != nullptr);
     return T(data_);
   }
+  /*!
+   * \return The internal object pointer with container type of T.
+   * \note This function do not perform not-null checking.
+   */
+  const ContainerType* get() const { return static_cast<ContainerType*>(data_.get()); }
   /*!
    * \return The contained value if the Optional is not null
    *         otherwise return the default_value.

diff --git a/src/tir/transforms/inject_rolling_buffer.cc b/src/tir/transforms/inject_rolling_buffer.cc
@@ -172,7 +172,7 @@ class RollingBufferInjector : public StmtExprMutator {
 
           auto it{std::find_if(
               bound_iter_vars.begin(), bound_iter_vars.end(),
-              [&](Optional<Var> var) { return var && (var.value().get() == loop_var.get()); })};
+              [&](Optional<Var> var) { return var && (var.get() == loop_var.get()); })};
 
           if (it != bound_iter_vars.end()) {
             auto i{std::distance(bound_iter_vars.begin(), it)};

diff --git a/src/tir/transforms/plan_update_buffer_allocation_location.cc b/src/tir/transforms/plan_update_buffer_allocation_location.cc
@@ -44,7 +44,7 @@ class BufferAllocationLocator : public StmtExprMutator {
     // create buffers to be allocated at each stmts
     for (const auto& kv : buffer_lca) {
       const Buffer& buffer = kv.first;
-      const StmtNode* stmt = kv.second.defined() ? kv.second.value().get() : nullptr;
+      const StmtNode* stmt = kv.second.get();
       if (arg_buffers.count(buffer.get())) {
         continue;
       }