[TIRx][CUDA] Configure kernel constraints with CUDA leaf ops

docs/tirx/native_basics/cuda/functions.rst

@@ -287,30 +287,36 @@ By default, the block size also drives the kernel's ``__launch_bounds__``. The
 first argument (max threads per block) is set automatically from the thread
 extent. To also set the second argument — the minimum blocks per SM, an
 occupancy hint — add
-``Tx.attr({"tirx.launch_bounds_min_blocks_per_sm": N})`` in the device region (note:
-``Tx.attr``, not ``func_attr``):
+``Tx.cuda.launch_bounds_min_blocks_per_sm(N)`` in the device region:
 
 .. code-block:: python
 
     Tx.device_entry()
-    Tx.attr({"tirx.launch_bounds_min_blocks_per_sm": 2})   # second launch-bounds arg
+    Tx.cuda.launch_bounds_min_blocks_per_sm(2)   # second launch-bounds arg
     bx = Tx.cta_id([1]); tx = Tx.thread_id([256])
     ...
 
 .. code-block:: c++
 
     extern "C" __global__ void __launch_bounds__(256, 2) scale_kernel(...) { ... }
 
-Without the attr the second argument is omitted (just ``__launch_bounds__(256)``).
+Without this declaration the second argument is omitted (just ``__launch_bounds__(256)``).
+
+``Tx.cuda.launch_bounds_max_blocks_per_cluster(N)`` supplies the third operand
+and requires a minimum-blocks declaration. ``Tx.cuda.max_registers_per_thread(N)``
+emits ``__maxnreg__(N)`` and cannot accompany launch bounds. These declarations
+accept positive integer constants. Matching repetitions are allowed; conflicting
+values are rejected. They configure the containing kernel and produce no runtime
+instructions at their textual position.
 
 Some kernels require an exact block and cluster shape instead of an advisory
-maximum. Set ``tirx.required_block_size`` to ``1`` to make the thread and
-cluster extents a compile-time launch contract:
+maximum. Use ``Tx.cuda.required_block_size`` with the three thread dimensions followed
+by the three cluster dimensions to declare a compile-time launch contract:
 
 .. code-block:: python
 
     Tx.device_entry()
-    Tx.attr({"tirx.required_block_size": 1})
+    Tx.cuda.required_block_size(128, 1, 1, 1, 2, 1)
     bx, by = Tx.cta_id([4, 2])
     _, cy = Tx.cta_id_in_cluster([1, 2])
     tx = Tx.thread_id([128])
@@ -321,15 +327,15 @@ cluster extents a compile-time launch contract:
     extern "C" __global__ void __block_size__((128, 1, 1), (1, 2, 1)) kernel(...) { ... }
 
 This requires CUDA Toolkit 13 or newer. All thread and cluster dimensions must
-be static; CUDA lowers ``__block_size__`` to PTX ``.reqntid`` and checks the
+be positive constants matching the declared launch extents; CUDA lowers ``__block_size__`` to PTX ``.reqntid`` and checks the
 same dimensions at launch. A preferred cluster dimension must be absent or
 equal to the required cluster dimension, and each logical block-grid dimension
 must be divisible by its cluster dimension.
 
-``tirx.required_block_size`` can be combined with the launch-bounds attributes
+``Tx.cuda.required_block_size`` can be combined with the launch-bounds declarations
 when an occupancy hint is also needed; code generation then emits both
 ``__block_size__`` and ``__launch_bounds__``. It cannot be combined with
-``tirx.max_registers``.
+``Tx.cuda.max_registers_per_thread``.
 
 At run time the kernel is launched through the **CUDA Driver API**. TVM's CUDA
 runtime loads the module (``cuModuleLoadData``), fetches the function
@@ -338,7 +344,7 @@ runtime loads the module (``cuModuleLoadData``), fetches the function
 the config carries a list of launch *attributes* — the thread-block **cluster
 dimension** and **preferred cluster dimension** (Hopper/Blackwell), plus optional
 programmatic-dependent-launch and cooperative-launch flags. Kernels with
-``tirx.required_block_size`` instead use CUDA's required-block sentinel; their
+``Tx.cuda.required_block_size`` instead use CUDA's required-block sentinel; their
 compile-time cluster shape replaces the ordinary runtime cluster attribute. In
 outline, ``src/backend/cuda/runtime/cuda_module.cc`` follows this path:
 

include/tvm/tirx/function.h

@@ -225,43 +225,6 @@ namespace attr {
  */
 constexpr const char* kKernelLaunchParams = "tirx.kernel_launch_params";
 
-/*!
- * \brief CUDA launch bound minimum CTAs per SM.
- *
- * Type: IntImm
- */
-constexpr const char* kLaunchBoundsMinBlocksPerSM = "tirx.launch_bounds_min_blocks_per_sm";
-
-/*!
- * \brief CUDA launch bound maximum CTAs per cluster.
- *
- * Type: IntImm
- */
-constexpr const char* kLaunchBoundsMaxBlocksPerCluster =
-    "tirx.launch_bounds_max_blocks_per_cluster";
-
-/*!
- * \brief CUDA maximum registers per thread.
- *
- * Emits the CUDA 13 ``__maxnreg__`` kernel qualifier.  This attribute is
- * mutually exclusive with the launch-bounds attributes.
- *
- * Type: IntImm
- */
-constexpr const char* kMaxRegisters = "tirx.max_registers";
-
-/*!
- * \brief Require CUDA to use the statically-declared block and cluster dimensions.
- *
- * Emits the CUDA 13 ``__block_size__`` kernel qualifier.  Unlike
- * ``__launch_bounds__``, this is an exact launch contract: CUDA derives the
- * PTX ``.reqntid`` directive from the thread extents, and interprets the
- * launch grid in clusters using the cluster-CTA extents.
- *
- * Type: IntImm (must be 1)
- */
-constexpr const char* kRequiredBlockSize = "tirx.required_block_size";
-
 /*!
  * \brief Whether to set noalias rule on the function arguments.
  *

src/backend/cuda/codegen/codegen_cuda.cc

@@ -175,6 +175,11 @@ CodeGenCUDA::CodeGenCUDA(Target target) : target(target) { restrict_keyword_ = "
 
 void CodeGenCUDA::PrintFunctionSignature(const ffi::String& function_name, const PrimFunc& func,
                                          std::ostream& os) {
+  PrintFunctionPrefix(func, os);
+  CodeGenC::PrintFunctionSignature(function_name, func, os);
+}
+
+void CodeGenCUDA::PrintFunctionPrefix(const PrimFunc& func, std::ostream& os) {
   CallingConv calling_conv =
       func->GetAttr<CallingConv>(tvm::attr::kCallingConv, CallingConv::kDefault).value();
   in_kernel_launch_ = (calling_conv == CallingConv::kDeviceKernelLaunch);
@@ -186,7 +191,6 @@ void CodeGenCUDA::PrintFunctionSignature(const ffi::String& function_name, const
     TVM_FFI_THROW(InternalError) << "Unsupported calling convention for cuda codegen: "
                                  << static_cast<int>(calling_conv);
   }
-  CodeGenC::PrintFunctionSignature(function_name, func, os);
 }
 
 class ThreadIdxExtractor : public tirx::StmtExprVisitor {
@@ -239,71 +243,98 @@ class ThreadIdxExtractor : public tirx::StmtExprVisitor {
   PrimExpr clusterCtaIdx_z_ext = IntImm::Int32(1);
 };
 
-void CodeGenCUDA::PrintExtraAttrs(const PrimFunc& f, std::ostream& os) {
+void CodeGenCUDA::InitFuncState(const PrimFunc& func) {
+  CodeGenC::InitFuncState(func);
+  min_blocks_per_sm_.reset();
+  max_blocks_per_cluster_.reset();
+  max_registers_per_thread_.reset();
+  required_block_size_.reset();
+  in_kernel_launch_ =
+      func->GetAttr<CallingConv>(tvm::attr::kCallingConv, CallingConv::kDefault).value() ==
+      CallingConv::kDeviceKernelLaunch;
+  // Thread metadata is also needed while emitting cluster index accesses.
   auto extractor = ffi::make_object<ThreadIdxExtractor>();
-  extractor->Visit(f->body);
+  extractor->Visit(func->body);
+  launch_dimensions_ = {extractor->threadIdx_x_ext,     extractor->threadIdx_y_ext,
+                        extractor->threadIdx_z_ext,     extractor->clusterCtaIdx_x_ext,
+                        extractor->clusterCtaIdx_y_ext, extractor->clusterCtaIdx_z_ext};
   sym::Analyzer analyzer;
-  PrimExpr threadIdx_ext = analyzer->Simplify(
-      extractor->threadIdx_x_ext * extractor->threadIdx_y_ext * extractor->threadIdx_z_ext);
-  PrimExpr cluster_cta_yz_ext =
-      analyzer->Simplify(extractor->clusterCtaIdx_y_ext * extractor->clusterCtaIdx_z_ext);
-  if (const IntImmNode* const cluster_cta_yz_ext_int = cluster_cta_yz_ext.as<IntImmNode>()) {
-    cluster_cta_x_is_linear_rank_ = cluster_cta_yz_ext_int->value == 1;
-  } else {
-    cluster_cta_x_is_linear_rank_ = false;
-  }
-  auto max_registers = f->GetAttr<int64_t>(tirx::attr::kMaxRegisters);
-  auto required_block_size = f->GetAttr<int64_t>(tirx::attr::kRequiredBlockSize);
-  if (required_block_size.has_value()) {
-    TVM_FFI_ICHECK_EQ(required_block_size.value(), 1);
-    TVM_FFI_ICHECK(!max_registers.has_value())
-        << tirx::attr::kRequiredBlockSize << " cannot be combined with maximum registers";
-    const auto* tx = extractor->threadIdx_x_ext.as<IntImmNode>();
-    const auto* ty = extractor->threadIdx_y_ext.as<IntImmNode>();
-    const auto* tz = extractor->threadIdx_z_ext.as<IntImmNode>();
-    const auto* cx = extractor->clusterCtaIdx_x_ext.as<IntImmNode>();
-    const auto* cy = extractor->clusterCtaIdx_y_ext.as<IntImmNode>();
-    const auto* cz = extractor->clusterCtaIdx_z_ext.as<IntImmNode>();
-    TVM_FFI_ICHECK(tx && ty && tz && cx && cy && cz)
-        << tirx::attr::kRequiredBlockSize << " requires static thread and cluster dimensions";
-    os << " __block_size__((" << tx->value << ", " << ty->value << ", " << tz->value << "), ("
-       << cx->value << ", " << cy->value << ", " << cz->value << "))";
-    if (!f->GetAttr<int64_t>(tirx::attr::kLaunchBoundsMinBlocksPerSM).has_value()) {
-      TVM_FFI_ICHECK(!f->GetAttr<int64_t>(tirx::attr::kLaunchBoundsMaxBlocksPerCluster).has_value())
-          << tirx::attr::kLaunchBoundsMaxBlocksPerCluster << " requires "
-          << tirx::attr::kLaunchBoundsMinBlocksPerSM;
-      return;
-    }
+  cluster_cta_x_is_linear_rank_ =
+      is_one(analyzer->Simplify(launch_dimensions_[4] * launch_dimensions_[5]));
+}
+
+void CodeGenCUDA::DeclareFunction(const GlobalVar& gvar, const PrimFunc& func) {
+  if (!RegisterFunctionName(gvar, func)) return;
+  // Definitions supply their body-derived qualifiers in AddFunction.
+  if (!func->body.has_value()) {
+    InitFuncState(func);
+    PrintFunctionSignature(GetFunctionName(gvar), func, fwd_decl_stream);
+    fwd_decl_stream << ";\n";
   }
-  if (max_registers.has_value()) {
-    TVM_FFI_ICHECK_GT(max_registers.value(), 0);
-    TVM_FFI_ICHECK(!f->GetAttr<int64_t>(tirx::attr::kLaunchBoundsMinBlocksPerSM).has_value() &&
-                   !f->GetAttr<int64_t>(tirx::attr::kLaunchBoundsMaxBlocksPerCluster).has_value())
-        << tirx::attr::kMaxRegisters << " cannot be combined with CUDA launch bounds";
-    os << " __maxnreg__(" << max_registers.value() << ")";
+}
+
+void CodeGenCUDA::AddFunction(const GlobalVar& gvar, const PrimFunc& func) {
+  DeclareFunction(gvar, func);
+  if (!func->body.has_value()) return;
+  InitFuncState(func);
+  std::ostringstream parameters;
+  PrintFunctionParameters(func, parameters);
+
+  // Generate the body once, retaining the settings declared by its leaf ops.
+  std::ostringstream previous_functions;
+  stream.swap(previous_functions);
+  stream << " {\n";
+  PreFunctionBody(func);
+  int func_scope = BeginScope();
+  PrintStmt(func->body.value());
+  EndScope(func_scope);
+  PrintIndent();
+  stream << "}\n\n";
+  std::string body = stream.str();
+  stream.swap(previous_functions);
+
+  std::ostringstream signature;
+  PrintFunctionPrefix(func, signature);
+  PrintType(func->ret_type, signature);
+  PrintExtraAttrs(func, signature);
+  signature << " " << GetFunctionName(gvar) << parameters.str();
+  fwd_decl_stream << signature.str() << ";\n";
+  stream << signature.str() << body;
+}
+
+void CodeGenCUDA::PrintExtraAttrs(const PrimFunc& f, std::ostream& os) {
+  TVM_FFI_ICHECK(!max_blocks_per_cluster_ || min_blocks_per_sm_)
+      << "CUDA maximum blocks per cluster requires minimum blocks per SM";
+  TVM_FFI_ICHECK(!max_registers_per_thread_ || (!min_blocks_per_sm_ && !max_blocks_per_cluster_))
+      << "Maximum registers per thread cannot be combined with CUDA launch bounds";
+  TVM_FFI_ICHECK(!max_registers_per_thread_ || !required_block_size_)
+      << "Required block size cannot be combined with maximum registers";
+  if (required_block_size_) {
+    for (size_t i = 0; i < launch_dimensions_.size(); ++i) {
+      const auto* dim = launch_dimensions_[i].as<IntImmNode>();
+      TVM_FFI_ICHECK(dim) << "Required block size requires static thread and cluster dimensions";
+      TVM_FFI_ICHECK_EQ(dim->value, (*required_block_size_)[i])
+          << "Required block size must agree with declared launch extents";
+    }
+    const auto& dims = *required_block_size_;
+    os << " __block_size__((" << dims[0] << ", " << dims[1] << ", " << dims[2] << "), (" << dims[3]
+       << ", " << dims[4] << ", " << dims[5] << "))";
+    if (!min_blocks_per_sm_) return;
+  }
+  if (max_registers_per_thread_) {
+    os << " __maxnreg__(" << *max_registers_per_thread_ << ")";
     return;
   }
-  if (const IntImmNode* const threadIdx_ext_int = threadIdx_ext.as<IntImmNode>()) {
-    if (threadIdx_ext_int->value == 1) {
-      // unable to extract the number of threads per block, hence directly return
-      return;
-    }
-    auto min_blocks_per_sm = f->GetAttr<int64_t>(tirx::attr::kLaunchBoundsMinBlocksPerSM);
-    auto max_blocks_per_cluster = f->GetAttr<int64_t>(tirx::attr::kLaunchBoundsMaxBlocksPerCluster);
-    if (min_blocks_per_sm.has_value()) {
-      TVM_FFI_ICHECK_GT(min_blocks_per_sm.value(), 0);
-      os << " __launch_bounds__(" << threadIdx_ext_int->value << ", " << min_blocks_per_sm.value();
-      if (max_blocks_per_cluster.has_value()) {
-        TVM_FFI_ICHECK_GT(max_blocks_per_cluster.value(), 0);
-        os << ", " << max_blocks_per_cluster.value();
-      }
-      os << ")";
-    } else {
-      TVM_FFI_ICHECK(!max_blocks_per_cluster.has_value())
-          << tirx::attr::kLaunchBoundsMaxBlocksPerCluster << " requires "
-          << tirx::attr::kLaunchBoundsMinBlocksPerSM;
-      os << " __launch_bounds__(" << threadIdx_ext_int->value << ")";
+  sym::Analyzer analyzer;
+  PrimExpr threads =
+      analyzer->Simplify(launch_dimensions_[0] * launch_dimensions_[1] * launch_dimensions_[2]);
+  if (const auto* count = threads.as<IntImmNode>(); count && count->value != 1) {
+    os << " __launch_bounds__(" << count->value;
+    if (min_blocks_per_sm_) {
+      os << ", " << *min_blocks_per_sm_;
+      if (max_blocks_per_cluster_) os << ", " << *max_blocks_per_cluster_;
     }
+    os << ")";
   }
 }
 
@@ -1684,6 +1715,39 @@ void CodeGenCUDA::DispatchAllocTensor(const BindNode* op, const CallNode* buffer
 }
 
 void CodeGenCUDA::Dispatch_(const EvaluateNode* op) {
+  if (const auto* call = op->value.as<CallNode>()) {
+    static const Op min_blocks = Op::Get("tirx.cuda.launch_bounds_min_blocks_per_sm");
+    static const Op max_blocks = Op::Get("tirx.cuda.launch_bounds_max_blocks_per_cluster");
+    static const Op max_registers = Op::Get("tirx.cuda.max_registers_per_thread");
+    static const Op required_block = Op::Get("tirx.cuda.required_block_size");
+    std::optional<int64_t>* setting = nullptr;
+    if (call->op.same_as(min_blocks)) {
+      setting = &min_blocks_per_sm_;
+    } else if (call->op.same_as(max_blocks)) {
+      setting = &max_blocks_per_cluster_;
+    } else if (call->op.same_as(max_registers)) {
+      setting = &max_registers_per_thread_;
+    }
+    if (setting) {
+      int64_t value = static_cast<int64_t>(call->args[0].as_or_throw<IntImm>()->value);
+      TVM_FFI_ICHECK_GT(value, 0) << call->op << " must be positive";
+      TVM_FFI_ICHECK(!setting->has_value() || setting->value() == value)
+          << "Conflicting " << call->op << " values";
+      *setting = value;
+      return;
+    }
+    if (call->op.same_as(required_block)) {
+      std::array<int64_t, 6> dimensions;
+      for (size_t i = 0; i < dimensions.size(); ++i) {
+        dimensions[i] = static_cast<int64_t>(call->args[i].as_or_throw<IntImm>()->value);
+        TVM_FFI_ICHECK_GT(dimensions[i], 0) << "Required block dimensions must be positive";
+      }
+      TVM_FFI_ICHECK(!required_block_size_ || *required_block_size_ == dimensions)
+          << "Conflicting required block size values";
+      required_block_size_ = dimensions;
+      return;
+    }
+  }
   if (auto value = op->value.as<PrimExpr>(); value && is_const_int(value.value())) return;
   CodeGenC::Dispatch_(op);
 }

src/backend/cuda/codegen/codegen_cuda.h

@@ -28,6 +28,8 @@
 #include <tvm/target/codegen.h>
 #include <tvm/tirx/op.h>
 
+#include <array>
+#include <optional>
 #include <string>
 #include <unordered_map>
 
@@ -48,6 +50,9 @@ class CodeGenCUDA final : public CodeGenC {
     });
   }
   // override behavior
+  void DeclareFunction(const GlobalVar& gvar, const PrimFunc& func) final;
+  void AddFunction(const GlobalVar& gvar, const PrimFunc& func) final;
+  void InitFuncState(const PrimFunc& func) final;
   void PrintFunctionSignature(const ffi::String& function_name, const PrimFunc& func,
                               std::ostream& os) final;
   void PrintExtraAttrs(const PrimFunc& f, std::ostream& os) final;  // NOLINT(*)
@@ -90,6 +95,14 @@ class CodeGenCUDA final : public CodeGenC {
                        bool skip_first_arg, std::ostream& os) final;  // NOLINT(*)
 
  private:
+  void PrintFunctionPrefix(const PrimFunc& func, std::ostream& os);
+  std::array<PrimExpr, 6> launch_dimensions_{IntImm::Int32(1), IntImm::Int32(1), IntImm::Int32(1),
+                                             IntImm::Int32(1), IntImm::Int32(1), IntImm::Int32(1)};
+  std::optional<int64_t> min_blocks_per_sm_;
+  std::optional<int64_t> max_blocks_per_cluster_;
+  std::optional<int64_t> max_registers_per_thread_;
+  std::optional<std::array<int64_t, 6>> required_block_size_;
+
   // Handle volatile loads
   void HandleVolatileLoads(const std::string& value, const TensorLoadNode* op,
                            std::ostream& os) final;

src/backend/cuda/op/target_builtin.cc

@@ -260,6 +260,21 @@ OpDef& RegisterDeviceIntrinsic(OpDef&& def, const char* op_namespace, CallEffect
 }
 
 void RegisterDeviceIntrinsics() {
+  // Kernel configuration declarations survive lowering until CUDA body generation.
+  RegisterDeviceIntrinsic(OpDef("tirx.cuda.launch_bounds_min_blocks_per_sm"), "cuda",
+                          CallEffectKind::kEmbedInfo, sig::arg<IntImm>("value"))
+      .set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void());
+  RegisterDeviceIntrinsic(OpDef("tirx.cuda.launch_bounds_max_blocks_per_cluster"), "cuda",
+                          CallEffectKind::kEmbedInfo, sig::arg<IntImm>("value"))
+      .set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void());
+  RegisterDeviceIntrinsic(OpDef("tirx.cuda.max_registers_per_thread"), "cuda",
+                          CallEffectKind::kEmbedInfo, sig::arg<IntImm>("value"))
+      .set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void());
+  RegisterDeviceIntrinsic(
+      OpDef("tirx.cuda.required_block_size"), "cuda", CallEffectKind::kEmbedInfo,
+      sig::arg<IntImm>("thread_x"), sig::arg<IntImm>("thread_y"), sig::arg<IntImm>("thread_z"),
+      sig::arg<IntImm>("cluster_x"), sig::arg<IntImm>("cluster_y"), sig::arg<IntImm>("cluster_z"))
+      .set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void());
   RegisterDeviceIntrinsic(OpDef("tirx.cuda.any_sync"), "cuda", CallEffectKind::kPure,
                           sig::arg<IntExpr>("mask"), sig::arg("pred"))
       .set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Int(32));