nvfuser-cu121-torch25 0.2.25.dev20250201__cp312-cp312-manylinux_2_28_x86_64.whl

nvfuser/include/nvfuser/serde/fusion_record.h

@@ -0,0 +1,124 @@
+// 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 <python_frontend/fusion_record.h>
+#include <serde/factory.h>
+
+namespace nvfuser::serde {
+
+// Forward definition for RecordFunctor
+struct RecordFunctor;
+
+// OpRecord Function Signatures
+// ========================================================================
+// Unary Functions
+typedef std::function<TensorView*(TensorView*)> unary_tv_fn;
+typedef std::function<Val*(Val*)> unary_val_fn;
+
+// ========================================================================
+// Binary Functions
+typedef std::function<TensorView*(TensorView*, TensorView*)> binary_tv_fn;
+typedef std::function<Val*(Val*, Val*)> binary_val_fn;
+typedef std::function<TensorView*(TensorView*, Val*)> binary_tv_val_fn;
+typedef std::function<TensorView*(Val*, TensorView*)> binary_val_tv_fn;
+
+// ========================================================================
+// Ternary Functions
+// Binary with Alpha Functions
+typedef std::function<TensorView*(TensorView*, TensorView*, TensorView*)>
+    ternary_tv_fn;
+typedef std::function<Val*(Val*, Val*, Val*)> ternary_val_fn;
+typedef std::function<TensorView*(TensorView*, TensorView*, Val*)>
+    ternary_tv_tv_val_fn;
+typedef std::function<TensorView*(TensorView*, Val*, TensorView*)>
+    ternary_tv_val_tv_fn;
+typedef std::function<TensorView*(Val*, TensorView*, TensorView*)>
+    ternary_val_tv_tv_fn;
+typedef std::function<TensorView*(Val*, Val*, TensorView*)>
+    ternary_val_val_tv_fn;
+typedef std::function<TensorView*(TensorView*, Val*, Val*)>
+    ternary_tv_val_val_fn;
+typedef std::function<TensorView*(Val*, TensorView*, Val*)>
+    ternary_val_tv_val_fn;
+
+// ========================================================================
+// Ternary with Alpha Functions
+typedef std::function<TensorView*(TensorView*, TensorView*, TensorView*, Val*)>
+    ternary_alpha_tv_fn;
+typedef std::function<Val*(Val*, Val*, Val*, Val*)> ternary_alpha_val_fn;
+typedef std::function<TensorView*(TensorView*, TensorView*, Val*, Val*)>
+    ternary_alpha_tv_tv_val_fn;
+typedef std::function<TensorView*(TensorView*, Val*, TensorView*, Val*)>
+    ternary_alpha_tv_val_tv_fn;
+typedef std::function<TensorView*(Val*, TensorView*, TensorView*, Val*)>
+    ternary_alpha_val_tv_tv_fn;
+typedef std::function<TensorView*(Val*, Val*, TensorView*, Val*)>
+    ternary_alpha_val_val_tv_fn;
+typedef std::function<TensorView*(TensorView*, Val*, Val*, Val*)>
+    ternary_alpha_tv_val_val_fn;
+typedef std::function<TensorView*(Val*, TensorView*, Val*, Val*)>
+    ternary_alpha_val_tv_val_fn;
+// ========================================================================
+
+//! The RecordFunctorFactory class is used to deserialize the flatbuffer
+//! RecordFunctor table. We create an enum type for each RecordFunctor class.
+//! Each template specialization has a unique RecordType and parser function.
+class RecordFunctorFactory
+    : public Factory<RecordFunctor, python_frontend::RecordFunctor*> {
+ public:
+  RecordFunctorFactory()
+      : Factory((nvfuser::toUnderlying(RecordType::MAX) + 1)) {
+    setupFunctionMaps();
+    registerAllParsers();
+  }
+
+ private:
+  void registerAllParsers();
+  void setupFunctionMaps();
+
+  // String to Operation maps
+  // Unary Functions
+  std::unordered_map<std::string, unary_tv_fn> unary_tv;
+  std::unordered_map<std::string, unary_val_fn> unary_val;
+
+  // Binary Functions
+  std::unordered_map<std::string, binary_tv_fn> binary_tv;
+  std::unordered_map<std::string, binary_val_fn> binary_val;
+  std::unordered_map<std::string, binary_tv_val_fn> binary_tv_val;
+  std::unordered_map<std::string, binary_val_tv_fn> binary_val_tv;
+
+  // Ternary Functions
+  // Binary with Alpha Functions
+  std::unordered_map<std::string, ternary_tv_fn> ternary_tv;
+  std::unordered_map<std::string, ternary_val_fn> ternary_val;
+  std::unordered_map<std::string, ternary_tv_tv_val_fn> ternary_tv_tv_val;
+  std::unordered_map<std::string, ternary_tv_val_tv_fn> ternary_tv_val_tv;
+  std::unordered_map<std::string, ternary_val_tv_tv_fn> ternary_val_tv_tv;
+  std::unordered_map<std::string, ternary_val_val_tv_fn> ternary_val_val_tv;
+  std::unordered_map<std::string, ternary_tv_val_val_fn> ternary_tv_val_val;
+  std::unordered_map<std::string, ternary_val_tv_val_fn> ternary_val_tv_val;
+
+  // Ternary with Alpha Functions
+  std::unordered_map<std::string, ternary_alpha_tv_fn> ternary_alpha_tv;
+  std::unordered_map<std::string, ternary_alpha_val_fn> ternary_alpha_val;
+  std::unordered_map<std::string, ternary_alpha_tv_tv_val_fn>
+      ternary_alpha_tv_tv_val;
+  std::unordered_map<std::string, ternary_alpha_tv_val_tv_fn>
+      ternary_alpha_tv_val_tv;
+  std::unordered_map<std::string, ternary_alpha_val_tv_tv_fn>
+      ternary_alpha_val_tv_tv;
+  std::unordered_map<std::string, ternary_alpha_val_val_tv_fn>
+      ternary_alpha_val_val_tv;
+  std::unordered_map<std::string, ternary_alpha_tv_val_val_fn>
+      ternary_alpha_tv_val_val;
+  std::unordered_map<std::string, ternary_alpha_val_tv_val_fn>
+      ternary_alpha_val_tv_val;
+};
+
+} // namespace nvfuser::serde

nvfuser/include/nvfuser/serde/polymorphic_value.h

@@ -0,0 +1,52 @@
+
+// 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 <runtime/executor_kernel_arg.h>
+#include <serde/factory.h>
+#include <serde/fusion_cache_generated.h>
+#include <visibility.h>
+#include <functional>
+#include <memory>
+
+namespace nvfuser::serde {
+
+//! The PolymorphicValueFactory class is used to deserialize the flatbuffer
+//! PolymorphicValue table. This factory creates Bool, ComplexDouble, Double,
+//! Long, CPU Scalar, and CUDA Tensor objects. These arguments are stored in
+//! KernelArgumentHolder, which is used to schedule the fusion in
+//! FusionKernelRuntime and to run a kernel in KernelExecutor.
+class PolymorphicValueFactory
+    : public Factory<PolymorphicValue, nvfuser::PolymorphicValue> {
+ public:
+  PolymorphicValueFactory()
+      : Factory((nvfuser::toUnderlying(PolymorphicValueData::MAX) + 1)) {
+    registerAllParsers();
+  }
+
+ private:
+  void registerAllParsers();
+};
+
+nvfuser::PolymorphicValue deserializePolymorphicValue(const Scalar* c);
+
+flatbuffers::Offset<PolymorphicValue> serializePolymorphicValue(
+    flatbuffers::FlatBufferBuilder& builder,
+    std::shared_ptr<nvfuser::PolymorphicValue> v);
+
+flatbuffers::Offset<Scalar> serializeScalarCpu(
+    flatbuffers::FlatBufferBuilder& builder,
+    const at::Tensor& tensor);
+
+NVF_API flatbuffers::Offset<Scalar> serializeScalar(
+    flatbuffers::FlatBufferBuilder& builder,
+    const nvfuser::PolymorphicValue& v,
+    nvfuser::DataType t);
+
+} // namespace nvfuser::serde

nvfuser/include/nvfuser/serde/utils.h

@@ -0,0 +1,34 @@
+// 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 <serde/fusion_cache_generated.h>
+#include <type.h>
+
+namespace nvfuser::serde {
+
+//! A function to map the serde dtype to its corresponding nvfuser prim dtype
+PrimDataType mapToNvfuserDtype(long data_type);
+
+//! A function to map the serde dtype to its corresponding nvfuser datatype
+nvfuser::DataType mapToDtypeStruct(long data_type);
+
+//! A function to map the serde dtype to its corresponding aten dtype
+at::ScalarType mapToAtenDtype(long data_type);
+
+template <typename T>
+std::vector<T> parseVector(const flatbuffers::Vector<T>* fb_vector) {
+  std::vector<T> result(fb_vector->begin(), fb_vector->end());
+  return result;
+}
+
+// Flatbuffer stores bool values as uint8_t.
+std::vector<bool> parseBoolVector(
+    const flatbuffers::Vector<uint8_t>* fb_vector);
+
+} // namespace nvfuser::serde

nvfuser/include/nvfuser/struct.inl

@@ -0,0 +1,127 @@
+// 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 <visibility.h>
+#include <functional>
+#include <memory>
+#include <string>
+
+namespace nvfuser {
+
+// Note [Struct Support in PolymorphicValue]
+//
+// PolymorphicValue supports structs, which is just a list of named fields. The
+// most straightforward way to support structs is to use a map from field name
+// to value, something like:
+//   template <typename T>
+//   using Struct = std::unordered_map<std::string, T>;
+//   using PolymorphicValue = DynamicType<Containers<Struct>, ...>;
+// However, the performance of this approach is not ideal. So instead of making
+// the struct support truly dynamic fields by using a map, we decide to make it
+// semi-dynamic: each struct type in nvFuser must be backed by a real struct in
+// C++, which mean, the fields have static storage types. But, on the other
+// hand, struct fields can also be accessed dynamically, that is, you can get or
+// set a struct field without knowing the actual C++ struct and the type of the
+// field. Instead, by using solely the string name of the field, you shall be
+// able to access fields as a PolymorphicValue. For example, if your struct is
+// defined as:
+//   struct A { int64_t x; double y; };
+//   PolymorphicValue v = some struct of type A;
+// Then you can access the fields statically like:
+//   const int64_t& x = v->*&A::x;
+//   v->*&A::x = 1;
+// Static accesses should be very efficient, as fast as dynamic casts + pointer
+// dereferences. However, if you don't have access to the definition of `A`, you
+// can still access the fields dynamically:
+//   PolymorphicValue x = v->*"x";
+//   v->*"x" = 1;
+// Dynamic accesses are slower than static accesses, because you need to do
+// string comparisons to find the field, and do casts between the actual field
+// type and PolymorphicValue. This can be slow especially when the struct has
+// some fields of containers like std::vector<int64_t>, because you need to do
+// the conversion between std::vector<PolymorphicValue> and std::vector<int64_t>
+// every time you get or set a field.
+//
+// The implementation of this feature requires a few components working
+// together:
+// 1. StructType: a data type that describes the name and fields of a struct.
+//    More importantly, it stores a function that can create an instance of a
+//    struct without requiring the caller to know the actual struct type.
+// 2. Struct: a base class for all structs, which provides the virtual interface
+//    for accessing fields dynamically, as well as an interface for getting the
+//    StructType of the struct.
+// 3. StructHandle: a wrapper around Struct, which maintains the ownership of
+//    struct objects and provides the overloaded ->* operator for accessing
+//    fields statically and dynamically. StructHandle is a candidate type for
+//    PolymorphicValue.
+// 4. Accessor: a helper class returned by the dynamic ->* operator, which
+//    provides the overloaded casting to PolymorphicValue and = operator for
+//    getting and setting fields dynamically.
+//
+// With the above components, define a struct type that supports dynamic access
+// to fields is basically subclassing Struct and implementing the virtual
+// methods. Please check the test PolymorphicValueTest.Struct for an example.
+
+struct Struct {
+  virtual ~Struct() = default;
+
+  virtual StructType type() const = 0;
+  virtual std::function<PolymorphicValue()> getter(
+      const std::string& key) const = 0;
+  virtual std::function<void(const PolymorphicValue&)> setter(
+      const std::string& key) = 0;
+};
+
+class Accessor {
+  std::function<PolymorphicValue()> getter_;
+  std::function<void(const PolymorphicValue&)> setter_;
+
+ public:
+  Accessor(
+      std::function<PolymorphicValue()> getter,
+      std::function<void(const PolymorphicValue&)> setter)
+      : getter_(std::move(getter)), setter_(std::move(setter)) {}
+  Accessor(const Accessor& value) = default;
+  Accessor(Accessor&& value) = default;
+  Accessor& operator=(const Accessor& value) = default;
+  Accessor& operator=(Accessor&& value) = default;
+
+  inline const Accessor& operator=(const PolymorphicValue& value) const {
+    setter_(std::move(value));
+    return *this;
+  }
+
+  inline operator PolymorphicValue() const {
+    return getter_();
+  }
+};
+
+inline Accessor StructHandle::operator->*(const std::string& key) const {
+  return Accessor(struct_ptr_->getter(key), struct_ptr_->setter(key));
+}
+
+// If a struct type is only used in kernel and we will never create an instance
+// on the host, we can just use this dummy struct as a placeholder for the
+// convenience
+struct NVF_API NotImplementedStruct : public Struct {
+  StructType type() const override;
+
+  std::function<PolymorphicValue()> getter(
+      const std::string& key) const override {
+    NVF_THROW("Not implemented");
+  }
+
+  std::function<void(const PolymorphicValue&)> setter(
+      const std::string& key) override {
+    NVF_THROW("Not implemented");
+  }
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/swizzle.h

@@ -0,0 +1,54 @@
+// 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 <visibility.h>
+
+#include <ir/interface_nodes.h>
+#include <type.h>
+
+#include <utility>
+
+//
+// The operations defined in this header is intended as user facing functions.
+// The user will provide the necessary input Vals and the function will
+// create the correct intermediate nodes and return the output Vals.
+//
+
+namespace nvfuser {
+
+NVF_API std::pair<Val*, Val*> dispatchSwizzle(
+    Swizzle2DType type,
+    Val* x,
+    Val* y,
+    Val* maybe_size_x,
+    Val* maybe_size_y);
+
+NVF_API std::pair<Val*, Val*> dispatchSwizzle(
+    SwizzleType type,
+    Val* x,
+    Val* y,
+    Val* maybe_size_x,
+    Val* maybe_size_y);
+
+NVF_API std::pair<Val*, Val*> dispatchUnSwizzle(
+    Swizzle2DType type,
+    Val* x,
+    Val* y,
+    Val* maybe_size_x,
+    Val* maybe_size_y);
+
+NVF_API std::pair<Val*, Val*> dispatchUnSwizzle(
+    SwizzleType type,
+    Val* x,
+    Val* y,
+    Val* maybe_size_x,
+    Val* maybe_size_y);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/sys_utils.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 <utils.h>
+
+#include <string>
+
+namespace nvfuser {
+
+class LibraryLoader : public NonCopyable {
+ public:
+  LibraryLoader() = default;
+
+  ~LibraryLoader();
+
+  void* getSymbol(const char* symbol_name);
+
+  std::string filename() const {
+    return filename_;
+  }
+
+  void setFilename(std::string filename) {
+    filename_ = filename;
+  }
+
+ private:
+  std::string filename_ = "";
+  void* handle_ = nullptr;
+};
+
+// Return true if compute-sanitizer is attached
+bool detectComputeSanitizer();
+
+} // namespace nvfuser

nvfuser/include/nvfuser/tensor_metadata.h

@@ -0,0 +1,118 @@
+// 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 <expr_evaluator.h>
+#include <ir/interface_nodes.h>
+#include <polymorphic_value.h>
+#include <type.h>
+
+namespace nvfuser {
+
+struct TensorMetaData : public Struct {
+  PrimDataType dtype;
+  void* data;
+  // References to the data fields. Does not own the data. The ownership might
+  // belong to at::Tensor or the *_data fields.
+  c10::IntArrayRef logical_size;
+  c10::IntArrayRef logical_stride;
+  c10::IntArrayRef alloc_size;
+  c10::IntArrayRef alloc_stride;
+  // The actual data for the above fields. Maybe empty if the fields are not
+  // owned by this object.
+  std::vector<int64_t> logical_size_data;
+  std::vector<int64_t> logical_stride_data;
+  std::vector<int64_t> alloc_size_data;
+  std::vector<int64_t> alloc_stride_data;
+
+  std::function<PolymorphicValue()> getter(
+      const std::string& key) const override {
+    if (key == "data") {
+      return [this]() { return PolymorphicValue(Pointer(data, dtype)); };
+    } else if (key == "logical_size") {
+      if (!logical_size_data.empty()) {
+        return [this]() { return PolymorphicValue(logical_size_data); };
+      } else {
+        return [this]() { return PolymorphicValue(logical_size.vec()); };
+      }
+    } else if (key == "logical_stride") {
+      if (!logical_stride_data.empty()) {
+        return [this]() { return PolymorphicValue(logical_stride_data); };
+      } else {
+        return [this]() { return PolymorphicValue(logical_stride.vec()); };
+      }
+    } else if (key == "alloc_size") {
+      if (!alloc_size_data.empty()) {
+        return [this]() { return PolymorphicValue(alloc_size_data); };
+      } else {
+        return [this]() { return PolymorphicValue(alloc_size.vec()); };
+      }
+    } else if (key == "alloc_stride") {
+      if (!alloc_stride_data.empty()) {
+        return [this]() { return PolymorphicValue(alloc_stride_data); };
+      } else {
+        return [this]() { return PolymorphicValue(alloc_stride.vec()); };
+      }
+    } else {
+      NVF_THROW("Unknown key ", key);
+    }
+  }
+
+  std::function<void(const PolymorphicValue&)> setter(
+      const std::string& key) override {
+    if (key == "data") {
+      return [this](const PolymorphicValue& value) { data = (void*)value; };
+    } else if (key == "logical_size") {
+      return [this](const PolymorphicValue& value) {
+        logical_size_data = (std::vector<int64_t>)value;
+        logical_size = c10::makeArrayRef(logical_size_data);
+      };
+    } else if (key == "logical_stride") {
+      return [this](const PolymorphicValue& value) {
+        logical_stride_data = (std::vector<int64_t>)value;
+        logical_stride = c10::makeArrayRef(logical_stride_data);
+      };
+    } else if (key == "alloc_size") {
+      return [this](const PolymorphicValue& value) {
+        alloc_size_data = (std::vector<int64_t>)value;
+        alloc_size = c10::makeArrayRef(alloc_size_data);
+      };
+    } else if (key == "alloc_stride") {
+      return [this](const PolymorphicValue& value) {
+        alloc_stride_data = (std::vector<int64_t>)value;
+        alloc_stride = c10::makeArrayRef(alloc_stride_data);
+      };
+    } else {
+      NVF_THROW("Unknown key ", key);
+    }
+  }
+
+  StructType type() const override {
+    NVF_ERROR(logical_size.size() == logical_stride.size());
+    NVF_ERROR(alloc_size.size() == alloc_stride.size());
+    return globalTensorMetaData(dtype, logical_size.size(), alloc_size.size());
+  }
+};
+
+// Given an ATen tensor, whose sizes and strides are w.r.t to the logical domain
+// of its corresponding TensorView, compute the sizes and strides of the tensor
+// with respect to its allocation domain.
+// For example, if the logical domain is [I1, I2], and the allocation domain is
+// [I2*I1], and the tensor's size is [5, 3] and stride is [2, 10], then the
+// resulting size will be [15] and stride will be [2]
+// Another example, if the logical domain is [I1*I2] and the allocation domain
+// is [I1, I2], and the tensor's size is [15] and stride is [7], and the extent
+// of I2 is 5, then the resulting size will be [3, 5] and stride will be [35, 7]
+std::pair<std::vector<int64_t>, std::vector<int64_t>>
+inferAndValidateAllocationSizesAndStrides(
+    const at::Tensor& tensor,
+    TensorView* tv,
+    ExpressionEvaluator ee);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/tma.h

@@ -0,0 +1,124 @@
+// 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 <ostream>
+
+#include <mma_type.h>
+#include <type.h>
+
+// Note: [TMA support in nvFuser]
+//
+// Recommended reading:
+// https://developer.nvidia.com/blog/nvidia-hopper-architecture-in-depth/
+// https://docs.nvidia.com/cuda/hopper-tuning-guide/index.html#tensor-memory-accelerator
+//
+// TMA (Tensor Memory Accelerator) is a hardware accelerator for transfering
+// tensors up to 5D between global memory and shared memory. It supports tiled
+// data and im2col data. nvFuser currently only supports tiled data.
+//
+// The tiled data transfer allows users to transfer a tile of a tensor between
+// global memory and shared memory. It is helpful to think of the tile as a
+// slice of the tensor. For example, if you have a 2D tensor you can do
+// something like:
+//   smem_tensor = gmem_tensor[i:i+16:2, j:j+32:1]
+// Or in the language of affine transformations, the gmem_tensor must be
+// transformed as:
+//   root domain: [I1, I2]
+//         split: [I1/16, 16, I2]
+//         split: [I1/16, 8, 2, I2]
+//         split: [I1/16, 8, 2, I2/32, 32]
+//   loop domain: [I1/16, 8, 2, I2/32, 32]
+//
+// Because TMA does bulk transfer, there is a dedicated paralle type `Bulk` for
+// it. In the above example, the gmem_tensor must be parallelized as
+//   [I1/16, Bulk{8}, 2, I2/32, Bulk{32}]
+// `Bulk` is a bit similar to `Vectorize` in some aspect, for example, both says
+// that we are copying a batch of data. Indeed, while considering `Bulk` as
+// representing a general N-dimensional slice that can have flexible extent and
+// step, we can consider `Vectorize` as a limited version of `Bulk` that must
+// represent a one-dimensional slice in the innermost dimension and the step of
+// the slice must be 1 and the extent of the slice must be a power of 2. Like
+// vectorize, a loop parallelized as `Bulk` is a trivial loop. Currently, we
+// only support whole tensor copy, so the consumer tensor of TMA store can not
+// be transformed, and all its `IterDomain`s must be parallelized as `Bulk`.
+//
+// To use TMA, we need to encode a tensor map of the global memory tensor we
+// want to transfer. The tensor map is a set of parameters that describes the
+// address and layout of the tensor in global memory, as well as the extent and
+// step of our slice. There are also other features configured in this tensor
+// map, such as: Does this tensor has overlap? How do we want to swizzle the
+// shared memory to avoid bank conflict? How do we want L2 cache to be used? Do
+// we want to automatically fill out-of-bound data? The tensor map must reside
+// in constant memory, and nvFuser implements this as a kernel argument declared
+// with __grid_constant__ qualifier. The tensor map is an opaque type to the
+// user. It must be created by driver API cuTensorMapEncodeTiled, see:
+// https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__TENSOR__MEMORY.html
+// In nvFuser, the task of calling `cuTensorMapEncodeTiled` to get the tensor
+// map is on the executor, implemented at: `kir::EncodeTensorMapTiled::evaluate`
+//
+// TMA are supported in nvFuser by inline PTX in memory.cu. Currently, we do a
+// sync immediately after the TMA PTX, which makes TMA synchronous. In the
+// future, we will need more advanced sync analysis and insert the correct syncs
+// at the correct point.
+//
+// During lowering, the index of the global tensor of the TMA expr must be
+// lowered as a kir::TensorIndex whose index has dtype `struct` with name
+// `Hopper::CpAsyncBulkTensorTileIndex`. The first field of this struct is the
+// pointer to the tensor map in constant memory, and the second field is an
+// array for the N-dimensional coordinate. The tensor map will be defined by an
+// expression of `kir::EncodeTensorMapTiled`. The evaluation of this expression
+// will be hoisted to the host, and will not be generated in the kernel. Because
+// we currently only support whole tensor copy, the N-dimensional coordinate is
+// always just zeros.
+//
+// Currently, because we only support very limited schedule, predicates for TMA
+// exprs are not needed, therefore not generated.
+
+namespace nvfuser {
+namespace tma {
+
+enum class TensorMapInterleave { NoInterleave, B16, B32 };
+enum class TensorMapL2Promotion { NoL2Promotion, B64, B128, B256 };
+enum class TensorMapFloatOOBFill { NoOOBFill, NaN_Request_Zero_FMA };
+
+std::ostream& operator<<(std::ostream& os, TensorMapInterleave interleave);
+std::ostream& operator<<(std::ostream& os, TensorMapL2Promotion l2_promotion);
+std::ostream& operator<<(std::ostream& os, TensorMapFloatOOBFill oob_fill);
+
+// Wrapper for:
+//
+// CUresult cuTensorMapEncodeTiled(
+//     CUtensorMap* tensorMap,
+//     CUtensorMapDataType tensorDataType,
+//     cuuint32_t tensorRank,
+//     void* globalAddress,
+//     const cuuint64_t* globalDim,
+//     const cuuint64_t* globalStrides,
+//     const cuuint32_t* boxDim,
+//     const cuuint32_t* elementStrides,
+//     CUtensorMapInterleave interleave,
+//     CUtensorMapSwizzle swizzle,
+//     CUtensorMapL2promotion l2Promotion,
+//     CUtensorMapFloatOOBfill oobFill);
+
+Val* encodeTensorMapTiled(
+    DataType data_type,
+    Val* global_address,
+    Val* global_dim,
+    Val* global_strides,
+    Val* box_dim,
+    Val* element_strides,
+    TensorMapInterleave interleave,
+    MmaInputSmemSwizzle swizzle,
+    TensorMapL2Promotion l2_promotion,
+    TensorMapFloatOOBFill oob_fill);
+
+} // namespace tma
+} // namespace nvfuser