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

nvfuser/include/nvfuser/scheduler/hopper_multi_matmul.h

@@ -0,0 +1,204 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <ATen/cuda/CUDAContext.h>
+#include <scheduler/multi_matmul.h>
+
+namespace nvfuser {
+
+// MmaOps in the scheduled tensor. Each one outputs a TensorView* which we call
+// an mma_result. Each MmaOp will also have two input TensorViews which we call
+// "ab" and "bb" since they are the immediate A and B operands and they contain
+// broadcast dimensions. Again there can be multiple abs and multiple bbs in
+// one fusion. These TensorViews are loaded from global memory tensors that we
+// call "a" and "b" into shared memory tensors called acw_smem and bcw_smem.
+// They are loaded from shared memory to register buffers we call "acr" and
+// "bcr" ("cr" meaning "cache read" in this context).
+//
+// Putting this all together we have the following order for a simple matmul
+//
+//   a -> acw_smem -> acr -> ... -> ab
+//                                    \                                      .
+//                                      mma_result ->  ... -> dc -> d
+//                                    /
+//   b -> bcw_smem -> bcr -> ... -> bb
+//
+// The ... indicate that there might be other tensors involved in a prologue or
+// epilogue section at that location.
+//
+// In this example there are two matmuls both using the same "a" operand:
+//
+//   b1 -> bcw_smem1 -> bcr1 -> ... -> bb1
+//                                        \                                  .
+//                                          mma_result1
+//                                        /             \                    .
+//       a -> acw_smem -> acr -> ... -> ab                ... -> dc -> d
+//                                        \             /
+//                                          mma_result2
+//                                        /
+//   b2 -> bcw_smem2 -> bcr2 -> ... -> bb2
+//
+// Note that there can be more than one output d and each one will have its own
+// register cache dc.
+//
+// Split-K and smem epilogue unswizzling add two additional tensors for each
+// mma in the fusion: splitk_sum and smem_epilogue.
+//
+//   // No split-K, no smem epilogue unswizzling:
+//     mma_result ->  ... -> dc -> d
+//   // split-K, no smem epilogue unswizzling:
+//     mma_result -> splitk_sum -> ... -> dc -> d
+//   // smem epilogue unswizzling, no split-K:
+//     mma_result -> smem_epilogue -> ... -> dc -> d
+//   // split-K and smem epilogue unswizzling:
+//     mma_result -> smem_epilogue -> splitk_sum -> ... -> dc -> d
+//
+// These additional tensors are added to each mma_result in the fusion.
+//
+// Each of the named tensors above is scheduled differently. We schedule them
+// by building AbstractTensors for each tensor category; these are held in
+// HopperMultipleMatmulScheduler::schedules_.
+// TODO: Inherit from SchedulerEntry
+class HopperMultipleMatmulScheduler : public MultipleMatmulScheduler {
+ public:
+  HopperMultipleMatmulScheduler(Fusion* fusion, const MatmulParams* params)
+      : MultipleMatmulScheduler(fusion, params) {
+    const auto device_prop = at::cuda::getCurrentDeviceProperties();
+    const int cc = device_prop->major * 10 + device_prop->minor;
+    NVF_ERROR(
+        cc >= 90 && cc < 100, "This matmul scheduler is restricted to Hopper.");
+  }
+
+  void run() final;
+
+ private:
+  void cacheInputsAndOutputs();
+
+  // Including current tensor naming convention for reference,
+  //  this is very temporary and will change over time and
+  //  in fact the whole body of this function will
+  //  eventually be a set of utility functions for different
+  //  sections of matmul(fusion) kernels, with
+  //  each having its own build out to do.
+  //
+  // Current naming convention is based on the following formula:
+  //
+  //  d = alpha * (a x b) + beta * c
+  //
+  // and is defined in the following way:
+  //
+  //  operands assumed in global memory : a, b, c
+  //
+  //  registers staging global load : ar, br (short for a/b read)
+  //
+  //  shared mem cache of operands : acw_smem, bcw_smem (short for a/b
+  //  cache_write smem)
+  //
+  //  registers at shared memory load output : acr, bcr (short for a/b cache
+  //  read)
+  //
+  //  register tensor input to the actual mma op: ab, bb (short for a/b
+  //  broadcasted)
+  //
+  //  accumulator register: mma_result
+  //   - mma_result is MmaOp output if there is epilogue
+  //   - mma_result is dc (short for d cache) if there is no epilogue
+  //
+  //  result in global memory: d
+
+  // Currently the support is for a, b, c and d as fusion inputs/outputs
+  //  aka. no prolog fusion yet.
+  void defineOperandCaches();
+
+  void cacheOperandsToSmem(
+      const std::vector<TensorView*>& operands,
+      std::vector<TensorView*>& smem_operands);
+
+  //! Swizzle the M and N outer dimensions after makeTile has been called.
+  //! This updates outer_dim_roles if we introduce a new dimension, which can
+  //! happen if tv is missing a merged axis, in which case we skip merging after
+  //! the split. This is analogous to forwarding during transform propagation.
+  void swizzleBlockTiles(
+      TensorView* tv,
+      std::vector<MatmulDimRole>& outer_dim_roles);
+
+  //! This calls orig->cacheAfter() and also updates the broadcast graph to
+  //! reflect the new IterDomain mappings
+  TensorView* cacheAfter(
+      TensorView* orig,
+      LoadStoreOpType op_type = LoadStoreOpType::Set,
+      CacheOp cache_op = CacheOp::AllLevels,
+      bool propagate_allocation_domain = false);
+
+  //! Do block tiling for a collection of TensorViews. The tensors should be
+  //! unscheduled before this method is called.
+  //!   1) Axes will be ordered according to canonicalDimOrdering, and then axes
+  //! with the same role will be merged.
+  //!   2) After that, we perform splits according to
+  //!   params_->tile_sizes.cta_tile, e.g. [M, K] -> [Mo, Ko, Mi, Ki].
+  //!   3) Depending on the value of params_->grid_swizzle_factor, if the TV has
+  //! both M and N dimensions, we perform a 2D swizzle of the outer dimensions
+  //! Mo and No.
+  //!   4) Finally, we do a split-K split if the splitk_factor is not 1
+  std::vector<std::vector<MatmulDimRole>> blockTileTensors(
+      const std::vector<TensorView*>& tvs);
+
+  //! Specifies the CGA dimensions by setting "cluster_dims" as fusion-managed
+  //! data
+  void setCGADims() const {
+    if (params_->cluster_dims != MatmulParams::ClusterDims{1, 1, 1}) {
+      fusion_->manage(
+          "cluster_dims",
+          std::tuple<int64_t, int64_t, int64_t>{
+              params_->cluster_dims.x,
+              params_->cluster_dims.y,
+              params_->cluster_dims.z});
+    }
+  }
+
+  //! Schedule the loads of all operands from global memory to shared memory.
+  //! Starting from the basic tiled schedule, we swizzle the operand memory.
+  //! Note that the cache op and LoadStoreOpType are already set during
+  //! defineOperandCaches().
+  void scheduleOperands();
+
+  //! Check that there is no computation in the prologues, since we do not
+  //! support that (yet)
+  void inspectPrologues() const;
+
+  void parallelizeBlocks(const std::vector<TensorView*>& tvs) const;
+
+  void scheduleMmaResults();
+
+  void scheduleEpilogue();
+
+  void scheduleSplitKSum();
+
+  void setUpInlining();
+
+  void setUpCircularBuffering();
+
+  // Map TensorView's iterDomain to its ValGroup.
+  // Then, find the MatmulDimRole for the ValGroup.
+  // Return MatmulDimRole for IterDomain
+  MatmulDimRole findMatmulDimRole(IterDomain* id);
+
+  // Schedule a block-tiled TensorView like mma output.
+  // Why? WGMMA has a unique output format. TensorViews after the mma-result in
+  // registers must respect this format for correctness.
+  void transformLikeMmaOutput(TensorView* tv, bool is_mma_result);
+
+ private:
+  std::vector<ValGroup> canonical_dim_ordering_;
+
+  std::vector<TensorView*> acw_smems_, bcw_smems_, acrs_, bcrs_, abs_, bbs_,
+      splitk_sums_, smem_epilogues_;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/scheduler/mark_aliases.h

@@ -0,0 +1,19 @@
+// 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>
+
+namespace nvfuser {
+
+// Marks aliases between fusion inputs and outputs. This respects existing
+// allocation domains **even when** they are empty (assuming default order). See
+// [Note on overriding empty allocation domains].
+void markAliases(Fusion* fusion);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/scheduler/matmul.h

@@ -0,0 +1,40 @@
+// 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/core/ivalue.h>
+#include <exceptions.h>
+#include <fusion.h>
+#include <mma_type.h>
+#include <scheduler/matmul_heuristic.h>
+#include <scheduler/registry.h>
+#include <visibility.h>
+
+namespace nvfuser {
+class MatmulScheduler : public SchedulerEntry {
+ public:
+  bool canScheduleCompileTime(Fusion* fusion) override;
+
+  bool canScheduleRunTime(
+      Fusion* fusion,
+      SchedulerRuntimeInfo& runtime_info,
+      HeuristicDataCache* data_cache = nullptr) override;
+
+  std::unique_ptr<HeuristicParams> computeHeuristics(
+      Fusion* fusion,
+      SchedulerRuntimeInfo& runtime_info,
+      HeuristicDataCache* data_cache = nullptr) override;
+
+  void schedule(Fusion* fusion, const HeuristicParams* params) override;
+
+  constexpr static SchedulerType schedulerType() {
+    return SchedulerType::Matmul;
+  }
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/scheduler/matmul_heuristic.h

@@ -0,0 +1,293 @@
+// 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 <c10/util/hash.h>
+#include <mma_type.h>
+#include <scheduler/heuristic.h>
+#include <utils.h>
+#include <functional>
+
+#include <sstream>
+#include "type.h"
+
+namespace nvfuser {
+
+// Parameters of the matmul heuristic to describe the optimial schedule.
+class MatmulParams : public HeuristicParams {
+ public:
+  MatmulParams()
+      : HeuristicParams(SchedulerType::Matmul), supported_vec_size() {};
+  //! A list of possible strategies used to define along which axis
+  //!  parallelization will be done.
+  enum class TileRasterizationOrder { RowMajor = 0, ColumnMajor = 1 };
+
+  //! A wrapper for circular buffering config pieces
+  struct CircularBufferOptions {
+    bool circular_buffer_smem_write = false;
+    bool circular_buffer_smem_read = false;
+    // This parameter controls the number of circular
+    // buffering stages to use when loading operands a and b.
+    //
+    // If this value is greater than two then it indicates circular buffering,
+    // in which case async_gmem_load_operands must also be true.
+    //
+    // Note that whenever circular_buffer_smem_write is true, this value must be
+    // greater than one. Otherwise it is ignored.
+    int smem_circular_buffer_stage = 2;
+
+    // The circular buffering prefetch distance will be set to
+    //   smem_circular_buffer_stage - smem_circular_buffer_prefetch_gap
+    // This value must be positive since the prefetch distance must be strictly
+    // less than the number of stages.
+    int smem_circular_buffer_prefetch_gap = 1;
+
+    bool operator==(const CircularBufferOptions& other) const {
+      return other.circular_buffer_smem_write == circular_buffer_smem_write &&
+          other.circular_buffer_smem_read == circular_buffer_smem_read &&
+          other.smem_circular_buffer_stage == smem_circular_buffer_stage &&
+          other.smem_circular_buffer_prefetch_gap ==
+          smem_circular_buffer_prefetch_gap;
+    }
+
+    std::string toString() const {
+      std::stringstream ss;
+      ss << "CircularBufferOptions:\n"
+         << "  circular_buffer_smem_write: "
+         << (circular_buffer_smem_write ? "true" : "false") << "\n"
+         << "  circular_buffer_smem_read: "
+         << (circular_buffer_smem_read ? "true" : "false") << "\n"
+         << "  smem_circular_buffer_stage: " << smem_circular_buffer_stage
+         << "\n"
+         << "  smem_circular_buffer_prefetch_gap: "
+         << smem_circular_buffer_prefetch_gap;
+      return ss.str();
+    }
+
+    size_t hash() const {
+      return std::hash<size_t>{}(
+                 (static_cast<size_t>(smem_circular_buffer_prefetch_gap) << 3) |
+                 (static_cast<size_t>(smem_circular_buffer_stage) << 2) |
+                 (static_cast<size_t>(circular_buffer_smem_write)) << 1) |
+          (static_cast<size_t>(circular_buffer_smem_read));
+    }
+  };
+
+  //! This is the maximum vectorization supported by the inputs and outputs.
+  //! This refers to the number of data elements loaded simultaneously, not the
+  //! number of bytes.
+  struct SupportedVectorization {
+    // Each operand load from global to shared memory is vectorized along its
+    // inner-most allocation dimension as long as that is an M, N, or K
+    // dimension. For example, if the innermost dimension is a batch dimension
+    // then we will not vectorize that operand's loads from global to shared
+    // memory. If there are multiple dimensions in a given role, such as
+    // multiple K dimensions, then we can only vectorize those inner dimensions
+    // that are consistent with the canonical dimension ordering shared by all
+    // tensors in the Fusion.
+    int64_t a;
+    int64_t b;
+
+    // The epilogue is handled in a separate loop from the main loop/operand
+    // loads. We inline the epilogue expressions as much as possible, and we
+    // vectorize all tensors with the same factor for better memory coalescence;
+    // i.e. we parallelize the epilogue like [ ... TIDx V ] so we do not
+    // introduce any loops between the TIDx and V dimensions. If we used
+    // different vectorization for each output or epilogue input, then we would
+    // need an unrolled loop between TIDx and V which would interfere with
+    // memory coalescence. We assume the decrease in indexing arithmetic from
+    // vectorization is not worth the slowdown from non-coalesced accesses, so
+    // we prefer to use a smaller vectorization instead.
+    //
+    // To determine the epilogue vectorization we do the following steps:
+    //  - Look at each output, then each epilogue input and find the first
+    //    tensor with a non-batch dimension as its innermost allocation
+    //    dimension. We will use that as the innermost loop dimension and will
+    //    vectorize that dimension. If there are multiple such innermost
+    //    dimensions with the same role and full contiguity then we consider all
+    //    those dimensions as the merged vectorized dimension. For example if
+    //    we have an output whose allocation domain is [ B1 M1 N1 M2 M3 ] then
+    //    (M2*M3) will be the vectorized dimension. On the other hand, we would
+    //    skip a tensor that had allocation domain [ M1 M2 M3 N1 B1 ] since the
+    //    batch dimension is innermost.
+    //  - Then we pass over all epilogue inputs and outputs. For each tensor, we
+    //    consider all innermost dimensions in order. For example if we have
+    //    determined that we will vectorize along M1*M2*M3 and a tensor has
+    //    allocation [ B1 M1 N1 M2 M3 ] then we consider dimension M2*M3 (along
+    //    with all other strides) to find supported vectorization. If another
+    //    tensor has allocation [ B1 M1 M2 M3 N1 ] then we skip it since its
+    //    innermost dimension is not an N role dimension so its access will not
+    //    be vectorized.
+    //  - We store the minimum of all the maximum supported vectorizations
+    //    across all epilogue input and output tensors that were not skipped.
+    //    That is the value below. If no vectorization is possible, this will be
+    //    set to 1.
+    int64_t epilogue;
+
+    bool operator==(const SupportedVectorization& other) const {
+      return other.a == a && other.b == b && other.epilogue == epilogue;
+    }
+
+    std::string toString() const {
+      std::stringstream ss;
+      ss << "SupportedVectorization:\n"
+         << "  a: " << a << "\n"
+         << "  b: " << b << "\n"
+         << "  epilogue: " << epilogue;
+      return ss.str();
+    }
+
+    size_t hash() const {
+      return std::hash<size_t>{}(
+                 (static_cast<size_t>(a) << 8) |
+                 (static_cast<size_t>(b)) << 4) |
+          (static_cast<size_t>(epilogue));
+    }
+  } supported_vec_size;
+
+  //! (Ampere+) Use cp.async to load operands.
+  bool async_gmem_load_operands = false;
+
+  //! Specifies the tiling hierarchy on block and warp levels.
+  MatMulTileOptions tile_sizes = {};
+
+  //! Specify the type of MMA op to be used in generated kernel.
+  MmaMacro mma_macro = MmaMacro::NoMMA;
+
+  //! Specify CTA rastrization order.
+  TileRasterizationOrder cta_order = TileRasterizationOrder::RowMajor;
+
+  //! Specify which tensor we circular buffer.
+  CircularBufferOptions circular_buffer_options = {};
+
+  //! Swizzle factor is used to increase L2 hit rate.
+  //!  It horizontally squeezes the grid so that gridDim.x is larger and
+  //!  gridDim.y is smaller.
+  //!  We rely on the observation that the CTAs are scheduled by the GPU by
+  //!  iterating on gridDim.x first. As a result, as blocks are launched, they
+  //!  will more likely be forming sub-tiles of the C matrix. This will increase
+  //!  L2 hit rate/data reuse of A and B.
+  //!
+  //! Eg for grid_swizzle_factor=2:
+  //!    A1 A2 B1 B2 -->   A1 A2 A3 A4 B1 B2 B3 B4
+  //!    A3 A4 B3 B4       C1 C2 C3 C4 D1 D2 D3 D4
+  //!    C1 C2 D1 D2
+  //!    C3 C4 D3 D4
+  int grid_swizzle_factor = 1;
+
+  //! Unswizzle MMA results in shared memory to get
+  //!  coalesced write to global memory
+  bool use_smem_epilogue = false;
+
+  //! Promote reuse of prologue shared memory
+  bool promote_prologue_smem_reuse = false;
+
+  //! Whether to do single-kernel split-K. If this is >1, we will rfactor the K
+  //! axis and perform a grid reduction before the epilogue.
+  int splitk_factor = 1;
+
+  //! This is the CGA size on Hopper+ devices. This parameter is ignored on
+  //! Ampere and Turing.
+  struct ClusterDims {
+    int64_t x = 1;
+    int64_t y = 1;
+    int64_t z = 1;
+
+    bool operator==(const ClusterDims& other) const {
+      return x == other.x && y == other.y && z == other.z;
+    }
+
+    bool operator!=(const ClusterDims& other) const {
+      return !(*this == other);
+    }
+
+    std::string toString() const {
+      std::stringstream ss;
+      ss << "__cluster_dims__(" << x << ", " << y << ", " << z << ")";
+      return ss.str();
+    }
+
+    size_t hash() const {
+      return std::hash<size_t>{}(
+                 (static_cast<size_t>(x) << 32) |
+                 (static_cast<size_t>(y)) << 16) |
+          (static_cast<size_t>(z));
+    }
+  } cluster_dims;
+
+  std::string toString() const override {
+    std::stringstream ss;
+    ss << "\n===== Matmul Parameters ========\n"
+       << (tag.empty() ? "" : "Tag: ") << tag << "\n"
+       << "MMA macro: " << nvfuser::toString(mma_macro) << "\n"
+       << circular_buffer_options.toString() << "\n"
+       << supported_vec_size.toString() << "\n"
+       << nvfuser::toString(tile_sizes) << "\n"
+       << "Async global mem load: "
+       << (async_gmem_load_operands ? "true" : "false") << "\n"
+       << "Indexing mode: "
+       << (cparams.index_type.has_value()
+               ? (cparams.index_type.value() == PrimDataType::Int ? "int64_t"
+                                                                  : "int32_t")
+               : "unavailable")
+       << "\n"
+       << "Tile rasterization order: "
+       << ((cta_order == TileRasterizationOrder::RowMajor) ? "row-major"
+                                                           : "column-major")
+       << "\n"
+       << "Grid swizzle factor: " << grid_swizzle_factor << "\n"
+       << cluster_dims.toString() << "\n"
+       << "Use shared memory epilogue: " << use_smem_epilogue << "\n"
+       << "Promote re-use of prologue shared memory: "
+       << promote_prologue_smem_reuse << "\n"
+       << "Split-K factor: " << splitk_factor << "\n"
+       << "====================================\n";
+    return ss.str();
+  }
+
+  size_t hash() const override {
+    // combine boolean flags for hashing
+    size_t attr_hash = (static_cast<size_t>(promote_prologue_smem_reuse) << 2) |
+        (static_cast<size_t>(use_smem_epilogue) << 1) |
+        (static_cast<size_t>(async_gmem_load_operands));
+
+    // combined hash
+    attr_hash = std::hash<size_t>{}(attr_hash) ^
+        (nvfuser::hash(mma_macro) << 1) ^
+        (circular_buffer_options.hash() << 2) ^
+        (nvfuser::hash(tile_sizes) << 3) ^
+        (std::hash<size_t>{}(static_cast<size_t>(cta_order)) << 4) ^
+        (std::hash<size_t>{}(grid_swizzle_factor) << 5) ^
+        (std::hash<size_t>{}(splitk_factor) << 6);
+    return attr_hash;
+  }
+
+  bool sameAs(const HeuristicParams* other_base) const override {
+    auto other = dynamic_cast<const MatmulParams*>(other_base);
+    if (other == nullptr) {
+      return false;
+    }
+
+    return other->cparams == cparams && other->mma_macro == mma_macro &&
+        other->async_gmem_load_operands == async_gmem_load_operands &&
+        other->tile_sizes == tile_sizes &&
+        other->circular_buffer_options == circular_buffer_options &&
+        other->supported_vec_size == supported_vec_size &&
+        other->cta_order == cta_order &&
+        other->grid_swizzle_factor == grid_swizzle_factor &&
+        other->use_smem_epilogue == use_smem_epilogue &&
+        other->promote_prologue_smem_reuse == promote_prologue_smem_reuse &&
+        other->splitk_factor == splitk_factor;
+  }
+
+  std::unique_ptr<HeuristicParams> clone() const override {
+    return std::make_unique<MatmulParams>(*this);
+  }
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/scheduler/matmul_heuristic_plugin.h

@@ -0,0 +1,65 @@
+// 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 <mma_type.h>
+#include <scheduler/matmul_heuristic.h>
+#include <scheduler/matmul_heuristic_plugin_api.h>
+#include <scheduler/mma_utils.h>
+
+#include <memory>
+#include <optional>
+
+namespace nvfuser {
+
+namespace matmul_heuristic_plugin {
+
+//! Returns true if KernelConfigFactoryGuard is active indicating an imitated
+//! plugin, or if a shared library plugin has been provided using the
+//! environment variable NVFUSER_MATMUL_HEURISTIC_PLUGIN.
+bool hasPlugin();
+
+//! If there is no user-defined plugin (see hasPlugin()) we return false.
+//! Otherwise, we use the plugin to modify the heuristic parameters in place. M,
+//! N, K, layout (inner allocated dimension roles of each operand), and
+//! precision must be provided. For convenience, we use `roles_map` to build the
+//! precision string.
+bool updateMatmulParams(
+    MatmulParams* params,
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t batch_size,
+    const mma_utils::MatmulOperandInnerDims& inner_dims,
+    const mma_utils::TensorRolesMap& tensor_roles);
+
+//! Defines the type of the "makeConfig" symbol
+using KernelConfigFactory = std::function<std::unique_ptr<KernelConfig>()>;
+
+//! This function can be used to imitate a plugin. To do so, subclass
+//! KernelConfig, implementing a custom `configure` method, then create a guard
+//! object like this:
+//!
+//!   KernelConfigFactoryGuard kfg([]() { return (KernelConfig*)(new
+//!   MyKernelConfig);});
+//!
+//! When kfg passes out of scope, the config factory will be reset to its prior
+//! value.
+class KernelConfigFactoryGuard {
+ public:
+  explicit KernelConfigFactoryGuard(KernelConfigFactory func);
+  ~KernelConfigFactoryGuard();
+
+ private:
+  KernelConfigFactory prev_factory_;
+  bool prev_factory_modified_;
+};
+
+} // namespace matmul_heuristic_plugin
+
+} // namespace nvfuser