Method UpdateCacheValue

oneflow/core/framework/op_kernel_infer_cache.cpp:59–77

}

void OpKernelInferCache::UpdateCacheValue(KernelInferContext* ctx) {
  // TODO: make max size configurable
  if (cached_key2value_.size() >= kReleaseInIndependentThreadThreshold) { Reset(); }
  auto* cache_value = new OpInferCacheValue();
  cache_value->obn_idx2shape_sym.resize(ctx->outputs().size());
  FOR_RANGE(int, i, 0, ctx->outputs().size()) {
    const auto& out_arg_pair = ctx->outputs().at(i);
    const ShapeView& out_shape_view =
        ctx->ShapeView4ArgNameAndIndex(out_arg_pair.first, out_arg_pair.second);
    Shape out_shape;
    out_shape_view.ToShape(&out_shape);
    cache_value->obn_idx2shape_sym.at(i).reset(out_shape);
  }
  KeyType* new_key = new KeyType(cache_key_);
  key_storage_.emplace_back(new_key);
  size_t hash_value = std::hash<KeyType>()(cache_key_);
  HashEqTraitPtr<const KeyType> ptr_wrapper(new_key, hash_value);
  CHECK(cached_key2value_.emplace(ptr_wrapper, ValueType(cache_value)).second);
}

void OpKernelInferCache::Reset() {
  CHECK_EQ(cached_key2value_.size(), key_storage_.size());