Expose MatmulParams::cluster_dims in python frontend

This introduces a new type `MatmulParams::ClusterDims` and wraps it in
the python frontend. This is essentially just a triple of `int64_t`s
but I named the elements x, y, and z to make it easy to bind.

Author: jacobhinkle

csrc/python_frontend/python_bindings.cpp

@@ -647,6 +647,13 @@ void defineHeuristicParamBindings(py::module& nvfuser) {
       nvfuser, "MatmulTileRasterizationOrder")
       .value("column_major", MatmulParams::TileRasterizationOrder::ColumnMajor)
       .value("row_major", MatmulParams::TileRasterizationOrder::RowMajor);
+
+  DEFINECLASS(MatmulParams::ClusterDims)
+      .PARAM(MatmulParams::ClusterDims, x)
+      .PARAM(MatmulParams::ClusterDims, y)
+      .PARAM(MatmulParams::ClusterDims, z)
+      .TOSTRINGMETHOD(MatmulParams::ClusterDims);
+
   py::enum_<MmaMacroEncode::Arch>(nvfuser, "MmaMacroArch")
       .value("no_mma", MmaMacroEncode::Arch::NoMma)
       .value("volta", MmaMacroEncode::Arch::Volta)
@@ -747,6 +754,7 @@ void defineHeuristicParamBindings(py::module& nvfuser) {
       .PARAM(MatmulParams, promote_prologue_smem_reuse)
       .PARAM(MatmulParams, splitk_factor)
       .PARAM(MatmulParams, cta_order)
+      .PARAM(MatmulParams, cluster_dims)
       .PARAM(MatmulParams, mma_macro);
 
 #undef PARAM

csrc/scheduler/hopper_multi_matmul.h

@@ -152,8 +152,13 @@ class HopperMultipleMatmulScheduler : public MultipleMatmulScheduler {
   //! Specifies the CGA dimensions by setting "cluster_dims" as fusion-managed
   //! data
   void setCGADims() const {
-    if (params_->cluster_dims != std::tuple<int, int, int>{1, 1, 1}) {
-      fusion_->manage("cluster_dims", params_->cluster_dims);
+    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});
     }
   }
 

csrc/scheduler/matmul_heuristic.h

@@ -193,7 +193,32 @@ class MatmulParams : public HeuristicParams {
 
   //! This is the CGA size on Hopper+ devices. This parameter is ignored on
   //! Ampere and Turing.
-  std::tuple<int64_t, int64_t, int64_t> cluster_dims = {1, 1, 1};
+  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 << 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;
@@ -216,8 +241,7 @@ class MatmulParams : public HeuristicParams {
                                                            : "column-major")
        << "\n"
        << "Grid swizzle factor: " << grid_swizzle_factor << "\n"
-       << "Cluster dimensions: " << std::get<0>(cluster_dims) << " "
-       << std::get<1>(cluster_dims) << " " << std::get<2>(cluster_dims) << "\n"
+       << "Cluster dimensions: " << cluster_dims.toString() << "\n"
        << "Use shared memory epilogue: " << use_smem_epilogue << "\n"
        << "Promote re-use of prologue shared memory: "
        << promote_prologue_smem_reuse << "\n"

csrc/scheduler/matmul_heuristic_plugin.cpp

@@ -139,9 +139,9 @@ void copyParamsToConfig(KernelConfig* config, const MatmulParams* mparams) {
   setConfigTile(config->cta_tile, mparams->tile_sizes.cta_tile);
   setConfigTile(config->warp_tile, mparams->tile_sizes.warp_tile);
   setConfigTile(config->instruction_tile, getMmaOpShape(mparams->mma_macro));
-  config->cluster_dims[0] = std::get<0>(mparams->cluster_dims);
-  config->cluster_dims[1] = std::get<1>(mparams->cluster_dims);
-  config->cluster_dims[2] = std::get<2>(mparams->cluster_dims);
+  config->cluster_dims[0] = mparams->cluster_dims.x;
+  config->cluster_dims[1] = mparams->cluster_dims.y;
+  config->cluster_dims[2] = mparams->cluster_dims.z;
   config->splitk_factor = mparams->splitk_factor;
   config->grid_swizzle_factor = mparams->grid_swizzle_factor;
   config->cta_order =
@@ -164,9 +164,9 @@ void copyConfigToParams(MatmulParams* mparams, const KernelConfig* config) {
   };
   setGemmTile(mparams->tile_sizes.cta_tile, config->cta_tile);
   setGemmTile(mparams->tile_sizes.warp_tile, config->warp_tile);
-  std::get<0>(mparams->cluster_dims) = config->cluster_dims[0];
-  std::get<1>(mparams->cluster_dims) = config->cluster_dims[1];
-  std::get<2>(mparams->cluster_dims) = config->cluster_dims[2];
+  mparams->cluster_dims.x = config->cluster_dims[0];
+  mparams->cluster_dims.y = config->cluster_dims[1];
+  mparams->cluster_dims.z = config->cluster_dims[2];
   mparams->circular_buffer_options.smem_circular_buffer_stage =
       config->load_stages;
   mparams->circular_buffer_options.smem_circular_buffer_prefetch_gap =