nvfuser-cu121-torch25 0.2.25.dev20250201__cp310-cp310-manylinux_2_28_x86_64.whl

nvfuser/include/nvfuser/global_allocator.h

@@ -0,0 +1,27 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <ATen/ATen.h>
+
+namespace nvfuser {
+
+//! This returns a slice of a thread local at::Tensor that contains all zeroes.
+//! Uses of this memory should always "clean up" by resetting the memory to zero
+//! at the end of the kernel.
+at::Tensor contigZeroedTensor(
+    const std::vector<int64_t>& sizes,
+    const c10::ScalarType& aten_dtype,
+    const c10::Device& device);
+
+//! This should be called after each kernel launch to allow subsequent launches
+//! to re-use allocated memory. Note that it does not free allocated zeroed
+//! memory, but rather it marks all zeroed memory as available for re-use.
+void releaseZeroedMemory();
+
+} // namespace nvfuser

nvfuser/include/nvfuser/grouped_reduction.h

@@ -0,0 +1,47 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <exceptions.h>
+#include <ir/all_nodes.h>
+#include <visibility.h>
+
+namespace nvfuser {
+
+//! Horizontally fuse multiple reductions.
+//!
+//! Given a list of tensors produced by ReductionOp, create a new
+//! GroupedReductionOp expression that takes the input tensors of the
+//! original reductions and produces the given tensors, replacing
+//! their defining expressions.
+//!
+//! GroupedReductionOp works just like ReductionOp with a potential
+//! benefit of aggregating synchronizations across individual
+//! reductions. See the reduction::gridReduce2 runtime function for a
+//! two-input version of grid reduction.
+//!
+//! The grouped reductions must follow several constraints, which
+//! include:
+//! - There must not exist any data dependency between individual
+//!   reductions.
+//! - All reduction output tensors must have the same number of
+//!   dimensions, the same transformations and the same axes to
+//!   reduce.
+//!
+//! Note that Welford is not allowed yet, though it should be
+//! technically straightforward to support horizontal fusions of
+//! welford ops. Unclear how common it would be in practice, though.
+//!
+//! \param reduction_outputs Tensors produced by ReductionOp
+//! \param error_on_failure Throw an exception if an error is detected
+//! \return True if successfully grouped
+NVF_API bool groupReductions(
+    const std::vector<TensorView*>& reduction_outputs,
+    bool error_on_failure = true);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/host_ir/container.h

@@ -0,0 +1,60 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <fusion.h>
+#include <host_ir/host_ir.h>
+
+namespace nvfuser {
+
+class KernelExecutor;
+
+namespace hir {
+
+/*
+HostIrContainer is used to represent a host program.
+1) It inherits from Fusion, so that (Host) IRs can be resgistered to it.
+2) It holds a vector of Host Expressions `top_level_exprs_` that represent the
+host program. For now, this vector is manually managed. Moreover, because we use
+a vector as data structure, top_level_exprs_ can only represent linear Host
+programs. Later, we it should support non-linear program having a DAG structure.
+*/
+
+class HostIrContainer final : public Fusion {
+ public:
+  HostIrContainer() = default;
+  HostIrContainer(const HostIrContainer&) = delete;
+  HostIrContainer& operator=(const HostIrContainer&) = delete;
+
+  // Do not have a definition here as it requires the definition of
+  // KernelExecutor due to kernel_executors_.
+  // NOLINTNEXTLINE (modernize-use-equals-default)
+  ~HostIrContainer() override;
+
+  //! Print to an output stream
+  std::ostream& print(std::ostream& os) const;
+
+  const std::vector<Expr*>& topLevelExprs() const;
+
+  void pushBackTopLevelExprs(Expr* expr);
+
+  void pushBackKernelExecutor(std::unique_ptr<KernelExecutor> ke);
+
+  KernelExecutor* getKernelExecutor(int64_t index) const;
+
+  Stream* getDefaultStream();
+
+ private:
+  std::vector<Expr*> top_level_exprs_;
+  std::vector<std::unique_ptr<KernelExecutor>> kernel_executors_;
+  Stream* default_stream_ = nullptr;
+};
+
+} // namespace hir
+
+} // namespace nvfuser

nvfuser/include/nvfuser/host_ir/executor.h

@@ -0,0 +1,152 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <dispatch.h>
+#include <expr_evaluator.h>
+#include <host_ir/container.h>
+#include <host_ir/host_ir.h>
+#include <multidevice/communicator.h>
+#include <runtime/executor.h>
+#include <runtime/executor_abstract.h>
+#include <runtime/executor_params.h>
+#include <runtime/fusion_executor_cache.h>
+
+#include <c10/cuda/CUDAStream.h>
+
+namespace nvfuser {
+
+class HostIrExecutor : public ExecutorAbstract {
+ public:
+  HostIrExecutor(
+      int64_t fusion_id = 0,
+      int64_t concrete_id = 0,
+      int64_t runtime_id = 0,
+      int64_t group_id = 0);
+
+  static bool supported(Fusion* fusion);
+
+  void compile(Fusion* fusion);
+
+  bool isCompiled() const override;
+
+  NVF_API std::vector<at::Tensor> run(
+      KernelArgumentHolder& args,
+      std::vector<at::Tensor> outputs = {});
+
+  const std::unique_ptr<hir::HostIrContainer>& hostContainer() const {
+    return host_ir_container_;
+  }
+
+ private:
+  std::unique_ptr<hir::HostIrContainer> host_ir_container_;
+  Communicator* communicator_;
+};
+
+namespace hir {
+
+/*
+a HostIrEvaluator evaluates a host programs represented through a
+HostIrContainer It is instantiated with the desired HostIrContainer, and runs
+the Host program with concrete inputs by calling the method runWithInput.
+
+For now HostIrEvaluator is an interpreter; later we could rather compile host
+code.
+
+Note: most of the implementation is copy pasted for MultiDeviceExecutor. This
+duplication will be resolved in the future.
+*/
+
+// Set of parameters that control the behavior of HostIrEvaluator
+struct HostIrEvaluatorParams {
+  // Experimental: whether to use FusionExecutorCache rather than
+  // KernelExecutor.
+  bool use_fusion_executor_cache = false;
+  // Experimental: whether to apply auto-scheduling in FusionExecutorCache if
+  // use_fusion_executor_cache=true. WAR: temporary hack mainly use for
+  // development
+  bool skip_auto_scheduling = false;
+  // Experimental: whether to cache fusion executor. WAR: avoid recompilation
+  // but implicitely assumes that the input shape don't change over iterations
+  bool cache_fusion_executor = false;
+  // number of additional cuda streams to use at runtime for comm+compute
+  // pipelining
+  int64_t number_of_streams = 4;
+};
+
+class HostIrEvaluator final : public OptOutDispatch {
+ public:
+  HostIrEvaluator(
+      std::unique_ptr<HostIrContainer> container,
+      Communicator* communicator = nullptr,
+      HostIrEvaluatorParams = HostIrEvaluatorParams());
+  std::vector<at::Tensor> runWithInput(
+      std::unordered_map<Val*, c10::IValue> val_to_IValue);
+
+  const std::vector<Val*>& inputs() {
+    return container_->inputs();
+  }
+
+  const std::vector<Val*>& outputs() {
+    return container_->outputs();
+  }
+
+  std::ostream& print(std::ostream& os) const {
+    return container_->print(os);
+  };
+
+  const auto& getFusionExecutorCaches() {
+    return fec_;
+  };
+
+  const auto& getCudaStreams() {
+    return streams_;
+  }
+
+  // check if the runtime is valid returns an error msg.
+  // An empty message means that the runtime is valid
+  std::string canRun() const;
+
+ private:
+  using OptOutDispatch::handle;
+  void handle(SetCurrentStream* set_current_stream) override;
+  void handle(GetCurrentStream* get_current_stream) override;
+  void handle(Synchronize* synchronize) override;
+  void handle(PostOnStream* post_ir) override;
+  void handle(LaunchKernel* post_ir) override;
+  void handle(Communication* communication) override;
+  void handle(P2PCommunication* communication) override;
+  void handle(Wait* wait) override;
+  void handle(ForLoop* for_loop) override;
+  void handle(StartCoalescing* start_coalescing) override;
+  void handle(EndCoalescing* end_coalescing) override;
+  void handle(kir::IfThenElse* if_then_else) override;
+  void handle(MatmulOp* matmul) override;
+  void handle(LinearOp* linear) override;
+  void handle(kir::Allocate* allocate) override;
+  void unhandled(Statement* stmt) override;
+
+  c10::cuda::CUDAStream getCUDAStream(Stream* stream);
+
+  std::unique_ptr<HostIrContainer> container_;
+  Communicator* communicator_;
+  HostIrEvaluatorParams params_;
+  // Stores concrete computed values
+  ExpressionEvaluator expr_evaluator_;
+  // Cache Fusions, KernelExecutors
+  std::unordered_map<HostUnit*, std::unique_ptr<ExecutorAbstract>> executors_;
+  std::unordered_map<HostUnit*, FusionExecutorCache> fec_;
+  using StreamKey = std::variant<int64_t, Stream*>;
+  std::unordered_map<StreamKey, c10::cuda::CUDAStream> streams_;
+  std::unordered_map<Expr*, c10::intrusive_ptr<c10d::Work>> works_;
+  const int64_t my_device_index_;
+};
+
+} // namespace hir
+
+} // namespace nvfuser

nvfuser/include/nvfuser/host_ir/host_ir.h

@@ -0,0 +1,320 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <fusion.h>
+#include <ir/base_nodes.h>
+#include <ir/builder.h>
+#include <multidevice/communication.h>
+#include <atomic>
+
+namespace nvfuser {
+
+namespace hir {
+
+/*
+  Host Irs are used to represent a host program. They need to be registered in a
+  HostIrContainer. Each Ir represents a Host data or instruction.
+*/
+
+/*
+  HostUnit represents a Fusion in the Host Program. In other words, it
+  represents a compute graph (or a segment of a larger compute graph)
+  represented by a Fusion that should be compiled and executed as a bulked item
+  from the host perspective.
+
+  This IR can be thought as a thin layer around the class `Fusion`, which
+  furthermore inherits from `Expr` so that it is an "IR" in nvFuser IR
+  semantics.
+
+  This IRs fundamentally allows nested IR structures. It could potentially be
+  useful in other instances than HostIrs.
+
+  Its implementation is minimal, the only specifity being the moethod
+  `fusion_to_execute()` that returns the fusion that the IR represents.
+
+  Note: HostUnit has no I/O itself -- however the Fusion it embbeds has I/O of
+  course, which are not registered in the surrounding HostIrContainer.
+
+  Note: Whether HostUnit should inherit from Expr or Val is debatable. Both are
+  possible, I define it as an Expr for now here but am open to change it.
+*/
+class HostUnit : public Expr {
+ public:
+  using Expr::Expr;
+  HostUnit(IrBuilderPasskey passkey, std::unique_ptr<Fusion> fusion);
+  HostUnit(const HostUnit* src, IrCloner* ir_cloner);
+
+  HostUnit(const HostUnit& other) = delete;
+  HostUnit& operator=(const HostUnit& other) = delete;
+  HostUnit(HostUnit&& other) = delete;
+  HostUnit& operator=(HostUnit&& other) = delete;
+
+  NVFUSER_DECLARE_CLONE_AND_CREATE
+  std::string toString(int indent_size = 0) const override;
+  std::string toInlineString(int indent_size = 0) const override;
+  const char* getOpString() const override {
+    return "hir::HostUnit";
+  }
+
+  bool sameAs(const Statement* other) const override;
+
+  Fusion* fusion_to_execute() const {
+    return fusion_.get();
+  }
+
+ private:
+  std::unique_ptr<Fusion> fusion_;
+};
+
+/*
+  PostOnStream represents the host instruction of executing a HostUnit. Its I/O
+  represents in the host program the concrete I/O that will be bound at runtime
+  to the Fusion's I/O for compilation and execution. At runtime, PostOnStream
+  will compile and launch the kernel lowered from the HostUnit's embedded
+  Fusion.
+
+  Note: later PostOnStream will take a "Stream" argument
+
+  Note: later PostOnStream will also be able to launch network Communications
+
+  Note: later compilation and kernel launch will be separated and represented by
+  distinct Host IRs
+*/
+class PostOnStream : public Expr {
+ public:
+  using Expr::Expr;
+  PostOnStream(
+      IrBuilderPasskey passkey,
+      Expr* host_op,
+      std::vector<Val*> inputs,
+      std::vector<Val*> outputs);
+
+  PostOnStream(const PostOnStream& other) = delete;
+  PostOnStream& operator=(const PostOnStream& other) = delete;
+  PostOnStream(PostOnStream&& other) = delete;
+  PostOnStream& operator=(PostOnStream&& other) = delete;
+
+  NVFUSER_DECLARE_CLONE_AND_CREATE
+
+  std::string toString(int indent_size = 0) const override;
+  std::string toInlineString(int indent_size = 0) const override;
+  const char* getOpString() const override {
+    return "hir::PostOnStream";
+  }
+
+  bool sameAs(const Statement* other) const override;
+
+  Expr* hostOpToPost() const {
+    return attributes_.at(0)->as<Expr>();
+  }
+};
+
+class LaunchKernel : public Expr {
+ public:
+  using Expr::Expr;
+  LaunchKernel(
+      IrBuilderPasskey passkey,
+      int64_t hic_executor_index, // Index into the HostIrContainer's vector of
+                                  // KernelExecutors--i.e., the kernel this IR
+                                  // should launch
+      const std::vector<Val*>& inputs,
+      const std::vector<Val*>& outputs);
+
+  LaunchKernel(const LaunchKernel& other) = delete;
+  LaunchKernel& operator=(const LaunchKernel& other) = delete;
+  LaunchKernel(LaunchKernel&& other) = delete;
+  LaunchKernel& operator=(LaunchKernel&& other) = delete;
+
+  NVFUSER_DECLARE_CLONE_AND_CREATE
+
+  std::string toString(int indent_size = 0) const override;
+  std::string toInlineString(int indent_size = 0) const override;
+  const char* getOpString() const override {
+    return "hir::LaunchKernel";
+  }
+
+  int64_t getIndex() const {
+    return attribute<int64_t>(0);
+  }
+};
+
+class Stream : public Val {
+ public:
+  // if index is provided, the IR represents the streams whose index is the
+  // dynamic value of that index. Otherwise, it statically represents a new
+  // Stream.
+  Stream(IrBuilderPasskey passkey, Val* index = nullptr);
+  Stream(const Stream* src, IrCloner* ir_cloner);
+  bool sameAs(const Statement* other) const override;
+
+  NVFUSER_DECLARE_CLONE
+  std::string toString(int indent_size = 0) const override;
+  std::string toInlineString(int indent_size = 0) const override;
+
+  Val* index() const {
+    return index_;
+  }
+
+ private:
+  Val* index_ = nullptr;
+};
+
+class SetCurrentStream : public Expr {
+ public:
+  using Expr::Expr;
+  SetCurrentStream(IrBuilderPasskey passkey, Stream* stream);
+
+  SetCurrentStream(const SetCurrentStream& other) = delete;
+  SetCurrentStream& operator=(const SetCurrentStream& other) = delete;
+  SetCurrentStream(SetCurrentStream&& other) = delete;
+  SetCurrentStream& operator=(SetCurrentStream&& other) = delete;
+
+  NVFUSER_DECLARE_CLONE_AND_CREATE
+
+  std::string toString(int indent_size = 0) const override;
+  std::string toInlineString(int indent_size = 0) const override;
+  const char* getOpString() const override {
+    return "hir::SetCurrentStream";
+  }
+
+  bool sameAs(const Statement* other) const override;
+
+  Stream* stream() const {
+    return attributes_.at(0)->as<Stream>();
+  }
+};
+
+class GetCurrentStream : public Expr {
+ public:
+  using Expr::Expr;
+  GetCurrentStream(IrBuilderPasskey passkey);
+
+  GetCurrentStream(const GetCurrentStream& other) = delete;
+  GetCurrentStream& operator=(const GetCurrentStream& other) = delete;
+  GetCurrentStream(GetCurrentStream&& other) = delete;
+  GetCurrentStream& operator=(GetCurrentStream&& other) = delete;
+
+  NVFUSER_DECLARE_CLONE_AND_CREATE
+
+  std::string toString(int indent_size = 0) const override;
+  const char* getOpString() const override {
+    return "hir::GetCurrentStream";
+  }
+
+  Stream* stream() const {
+    return attributes_.at(0)->as<Stream>();
+  }
+};
+
+class Wait : public Expr {
+ public:
+  using Expr::Expr;
+  Wait(IrBuilderPasskey passkey, Expr* expr);
+
+  Wait(const Wait& other) = delete;
+  Wait& operator=(const Wait& other) = delete;
+  Wait(Wait&& other) = delete;
+  Wait& operator=(Wait&& other) = delete;
+
+  NVFUSER_DECLARE_CLONE_AND_CREATE
+
+  std::string toString(int indent_size = 0) const override;
+  std::string toInlineString(int indent_size = 0) const override;
+  const char* getOpString() const override {
+    return "hir::Wait";
+  }
+
+  bool sameAs(const Statement* other) const override;
+
+  Expr* communication() const {
+    return attributes_.at(0)->as<Expr>();
+  }
+};
+
+// Makes the current stream wait on the given stream. Non-blocking from the host
+// point of view.
+class Synchronize : public Expr {
+ public:
+  using Expr::Expr;
+  Synchronize(IrBuilderPasskey passkey, Stream* stream);
+
+  Synchronize(const Synchronize& other) = delete;
+  Synchronize& operator=(const Synchronize& other) = delete;
+  Synchronize(Synchronize&& other) = delete;
+  Synchronize& operator=(Synchronize&& other) = delete;
+
+  NVFUSER_DECLARE_CLONE_AND_CREATE
+
+  std::string toString(int indent_size = 0) const override;
+  std::string toInlineString(int indent_size = 0) const override;
+  const char* getOpString() const override {
+    return "hir::Synchronize";
+  }
+
+  bool sameAs(const Statement* other) const override;
+
+  Stream* stream() const {
+    return attributes_.at(0)->as<Stream>();
+  }
+};
+
+// For ProcessGroupNCCL, startCoalescing and endCoalescing correspond to
+// ncclGroupStart and ncclGroupEnd respectively. Those calls group p2p calls
+// that need to be progressed together -- one global work handle returned by
+// endCoalescing needs to be progressed. This has the following main advantages:
+// 1) calls are progressed concurrently
+// 2) since NICs are two-sided, a send and a recv calls need to be coalesced to
+//    achieve full BW.
+// 3) If not coalesced, we can easily reach a deadlock if the
+//    send/recv pairs are not ordered correctly.
+// It is in general preferable to coalesce send/recv calls. The only drawback is
+// that we don't have a fine-grain control on synchronicity, in other words, we
+// can only synchronize with the grouped communication at once.
+// Remark: ProcessGroupUCC does not implement coalesced groups for now
+class StartCoalescing : public Expr {
+ public:
+  using Expr::Expr;
+  StartCoalescing(IrBuilderPasskey passkey);
+
+  StartCoalescing(const StartCoalescing& other) = delete;
+  StartCoalescing& operator=(const StartCoalescing& other) = delete;
+  StartCoalescing(StartCoalescing&& other) = delete;
+  StartCoalescing& operator=(StartCoalescing&& other) = delete;
+
+  NVFUSER_DECLARE_CLONE_AND_CREATE
+
+  std::string toString(int indent_size = 0) const override;
+  std::string toInlineString(int indent_size = 0) const override;
+  const char* getOpString() const override {
+    return "hir::StartCoalescing";
+  }
+};
+
+class EndCoalescing : public Expr {
+ public:
+  using Expr::Expr;
+  EndCoalescing(IrBuilderPasskey passkey);
+
+  EndCoalescing(const EndCoalescing& other) = delete;
+  EndCoalescing& operator=(const EndCoalescing& other) = delete;
+  EndCoalescing(EndCoalescing&& other) = delete;
+  EndCoalescing& operator=(EndCoalescing&& other) = delete;
+
+  NVFUSER_DECLARE_CLONE_AND_CREATE
+
+  std::string toString(int indent_size = 0) const override;
+  std::string toInlineString(int indent_size = 0) const override;
+  const char* getOpString() const override {
+    return "hir::EndCoalescing";
+  }
+};
+
+} // namespace hir
+
+} // namespace nvfuser

nvfuser/include/nvfuser/host_ir/lower.h

@@ -0,0 +1,35 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <host_ir/container.h>
+#include <ir/base_nodes.h>
+#include <multidevice/communication.h>
+#include <multidevice/multidevice.h>
+
+namespace nvfuser {
+
+class HostIrLower {
+ public:
+  // The flag `ignore_inner_resharding` is useful because the preseg passes
+  // `InsertReshardingsPass` and `ReorderShardedAxisPass` want different
+  // behaviors
+  static bool canLower(Expr* expr, bool ignore_inner_resharding = false);
+
+  // Lower a sharded Expr into a series of Communication.
+  static std::vector<Expr*> lower(Expr* c);
+
+  static std::unique_ptr<hir::HostIrContainer> lower(
+      std::unique_ptr<Fusion> fusion,
+      int64_t my_device_index);
+
+ private:
+  static std::vector<Expr*> lowerToCollectiveBasedPipelinedGemmComm(Expr* expr);
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/id_model/circular_buffer_indexing.h

@@ -0,0 +1,56 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <device_lower/lower2device.h>
+#include <device_lower/utils.h>
+#include <id_model/id_model.h>
+
+namespace nvfuser {
+
+// Get the loop index of a given loop domain for circular buffer
+// loops. nullptr is returned if not relevant.
+//
+// This is a WAR for circular buffering. TensorIndexer has a map of
+// loop indices for all loop groups, however, it does not work with
+// circular buffering. The loop graph is
+// designed to represent each loop and each loop group is supposed
+// to have a one-to-one relationship with each loop. However, for
+// circular buffering, this assumption is broken as we are using
+// the same iter domain for the prologue, main and epilogue
+// loops. Ideally, those loops should have distinctive loop groups,
+// but for now, here's a workaround to get a correct loop index
+Val* getLoopIndexOfCircularBufferLoop(
+    IterDomain* loop_id,
+    const std::vector<ForLoop*>& for_loops,
+    const IdModel& id_model);
+
+// For a circular-buffering expr, the producer loop index needs to be
+// advanced by (#stages - 1) if it's the main loop. Return the offset
+// if it's applicable. Otherwise, nullptr is returned.
+Val* getLoopIndexOffsetForProducerOfCircularBuffer(
+    const Expr* expr,
+    const ForLoop* for_loop,
+    const IdModel& id_model);
+
+// Get the additional offset for a circular buffer. This offset will
+// be added to the normal linear index. For example, if this is a
+// double buffered tensor, the offset would look like "i % 2", where i
+// is the loop index of the double-buffer loop.
+Val* getOffsetForCircularBufferTensor(
+    TensorView* circular_buffer_tv,
+    bool as_consumer,
+    const std::vector<ForLoop*>& for_loops);
+
+// Find the circular buffering stage of a given circular buffered tensor
+CircularBufferLoopStage getCircularBufferLoopStage(
+    const TensorView* circular_buffer_tv,
+    const std::vector<ForLoop*>& for_loops,
+    const ValGraph& loop_graph);
+
+} // namespace nvfuser