Add a script to run a single matmul configuration with custom MatmulParams

[rdspring1]

This PR adds a script to run a single matmul configuration with custom MatmulParams. It profiles the nvfuser kernel and compares its runtime against nvjet kernel runtime. The nvjet kernel runtimes are stored in a json file, generated by python benchmarks.

• Update python bindings to support constructing Matmul Options.

[rdspring1]

!test

[github-actions]

Review updated until commit f269852

Description

• Add script for profiling single matmul configuration

• Update Python bindings for MatmulParams

• Implement custom scheduler for matmul

• Validate nvFuser against PyTorch matmul

---

Changes walkthrough 📝

Relevant files

Enhancement

python_bindings.cpp Update Python bindings for MatmulParams                                    csrc/python_frontend/python_bindings.cpp Add constructors for GemmTile, MatMulTileOptions, CircularBufferOptions, SupportedVectorization, ClusterDims, and MmaMacroEncode Update class definitions to use py::class_ instead of DEFINECLASS macro +23/-4    profile_matmul.py Add script for profiling single matmul configuration          doc/dev/python_scheduling/profile_matmul.py Add script to run and profile a single matmul configuration Implement functions for estimating matmul size, getting kernel time, and defining matmul fusion Implement custom scheduler with custom parameters Add main function to parse arguments and run profiling +209/-0

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review Initialization Order The order of initialization for MatmulParams classes should be consistent with the class definitions to avoid potential issues. .def(py::init<int64_t, int64_t, int64_t>()) .PARAM(GemmTile, m) .PARAM(GemmTile, n) .PARAM(GemmTile, k) .TOSTRINGTOPLEVEL(GemmTile); DEFINECLASS(MatMulTileOptions) .def(py::init<GemmTile, GemmTile>()) .PARAM(MatMulTileOptions, cta_tile) .PARAM(MatMulTileOptions, warp_tile) .TOSTRINGTOPLEVEL(MatMulTileOptions); py::class_<MatmulParams::CircularBufferOptions>( Missing Definitions The DEFINECLASS macro is used instead of py::class_ for some classes, which might lead to missing bindings or incorrect behavior. .def(py::init<int64_t, int64_t, int64_t>()) .PARAM(GemmTile, m) .PARAM(GemmTile, n) .PARAM(GemmTile, k) .TOSTRINGTOPLEVEL(GemmTile); DEFINECLASS(MatMulTileOptions) .def(py::init<GemmTile, GemmTile>()) .PARAM(MatMulTileOptions, cta_tile) .PARAM(MatMulTileOptions, warp_tile) .TOSTRINGTOPLEVEL(MatMulTileOptions); py::class_<MatmulParams::CircularBufferOptions>( Error Handling The error handling in test_matmul_nvf could be improved to provide more informative messages or handle specific exceptions. nvf_outputs = scheduled_fd.execute([a, b], profile=True) except Exception as e: if verbose: print(e) return -1

[jacobhinkle]

Looks good! Thanks for adding this. Comments are pretty minor I think.

[jacobhinkle]

This is the bare minimum size required to compute the GEMM. If there is validation enabled, there will be more than that because of more outputs: 1 for the eager output plus intermediates required for torch.allclose. So you might want to add a fudge factor.

[rdspring1]

@jacobhinkle I added the following argparse.

usage: profile_matmul.py [-h] [--verbose] [--validate] m n k {NN,NT,TN,TT}

Run through a combination of matmul parameters and compare relative performance against nvjet for a single problem.

positional arguments:
  m              The size of M dimension
  n              The size of N dimension
  k              The size of K dimension
  {NN,NT,TN,TT}  The layout for matmul problem.

options:
  -h, --help     show this help message and exit
  --verbose      Print matmul parameters and exceptions.
  --validate     Validate nvfuser against pytorch matmul.

How to run script: NVFUSER_ENABLE=fuse_matmul NVFUSER_DISABLE=matmul_expr_eval python single_matmul.py nvjet_pybench.json 1752 4720 584 NN --verbose --validate

[rdspring1]

!build