Set __launch_bounds__ in kernel whenever we are able #3794

jacobhinkle · 2025-01-29T18:18:00Z

Currently we set the number of threads per block via __launch_bounds__ when register sharing is enabled. This PR just enables this whenever it is possible, i.e. whenever we know the CTA size at compile time.

Adds the method ParallelDimensionMap::getNumThreadsEachBlock() which is similar to ParallelDimensionMap::getNumComputeThreadsEachBlock() but includes all threads and doesn't skip dma threads.

See https://docs.nvidia.com/cuda/cuda-c-programming-guide/#launch-bounds for more background.

Currently we set the number of threads per block via `__launch_bounds__` when register sharing is enabled. This PR just enables this whenever it is possible, i.e. whenever we know the CTA size at compile time. See https://docs.nvidia.com/cuda/cuda-c-programming-guide/#launch-bounds for more background.

jacobhinkle · 2025-01-29T18:18:06Z

!test

github-actions · 2025-01-29T18:19:36Z

PR Reviewer Guide 🔍

(Review updated until commit `b78ae42`)

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 2 🔵🔵⚪⚪⚪
🧪 No relevant tests
⚡ Recommended focus areas for review Const Cast The code uses a const_cast to avoid a const correctness issue. Consider re-designing the code to avoid the need for const_cast. // Avoid a const_cast that would be required to use kernel_ by picking the // fusion of the first kernel output FusionGuard fg(kernel_->outputs().front()->fusion()); Thread Count Calculation The getNumThreadsEachBlock method calculates the total number of threads per block by multiplying the number of threads for each parallel type. Verify that this calculation is correct and accounts for all possible parallel types. Val* ParallelDimensionMap::getNumThreadsEachBlock() const { Val* num_threads = FusionGuard::getCurFusion()->oneVal(); for (auto pt : kParallelTypeTIDs) { num_threads = SimplifyingIrBuilder::mulExpr(num_threads, getRaw(pt)); } return num_threads; }

zasdfgbnm

Makes sense to me, but will let Ryan to decide.

rdspring1 · 2025-01-29T18:33:09Z

!test --pybench

rdspring1

num_threads_per_cta should always have a value because it is the inferred launch bounds. https://github.com/NVIDIA/Fuser/blob/main/csrc/runtime/executor.cpp#L268-L286

This seems different from what you intended. I thought you wanted to set the launch bounds when TIDx, TIDy, and TIDz extents are constant.

jacobhinkle · 2025-01-29T20:59:30Z

This seems different from what you intended. I thought you wanted to set the launch bounds when TIDx, TIDy, and TIDz extents are constant.

Ah, you're right. I did intend it to not be derived from inputs since that would interfere with dynamic shapes. I'll give it another try.

jacobhinkle · 2025-01-30T20:21:17Z

!test --diff

jacobhinkle · 2025-01-31T01:24:37Z

!test --diff

jacobhinkle · 2025-01-31T13:12:17Z

!test --diff

github-actions · 2025-01-31T13:12:53Z

Review updated until commit 0088161

Description

Set __launch_bounds__ whenever CTA size is known at compile time.
Add ParallelDimensionMap::getNumThreadsEachBlock() method.
Update tests to reflect new kernel launch bounds.

Changes walkthrough 📝

Relevant files

Enhancement

codegen.cpp `Add logic for setting launch bounds` csrc/codegen.cpp Added logic to evaluate block size and set `__launch_bounds__` if CTA size is known. Updated kernel declaration to include `__launch_bounds__` when available.	+12/-0
parallel_dimension_map.cpp `Implement getNumThreadsEachBlock method` csrc/parallel_dimension_map.cpp Implemented `getNumThreadsEachBlock()` to calculate total threads per block.	+8/-0
parallel_dimension_map.h `Add getNumThreadsEachBlock declaration` csrc/parallel_dimension_map.h Added declaration for `getNumThreadsEachBlock()`.	+3/-0

Tests

test_loop_rotation.cpp `Update expected kernel strings` tests/cpp/test_loop_rotation.cpp Updated expected kernel strings to include `__launch_bounds__`.	+6/-6
test_scalar_hoisting.cpp `Update expected kernel strings` tests/cpp/test_scalar_hoisting.cpp Updated expected kernel strings to include `__launch_bounds__`.	+2/-2

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review

Launch Bounds Calculation

Ensure that the calculation of num_threads_per_cta is correct and handles all edge cases, especially when num_threads_per_cta is not initially provided.

if (!num_threads_per_cta.has_value()) {
  // Try to evaluate the block size so that we can set launch bounds.
  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
  FusionGuard fg(const_cast<kir::Kernel*>(kernel_));
  Val* num_threads =
      kernel_->summary().parallel_dimension_map.getNumThreadsEachBlock();
  if (num_threads->isConstInt()) {
    num_threads_per_cta = num_threads->evaluate().as<int64_t>();
  }
}

Thread Count Calculation

Verify that getNumThreadsEachBlock correctly calculates the total number of threads per block, including all types of threads.

Val* ParallelDimensionMap::getNumThreadsEachBlock() const {
  Val* num_threads = FusionGuard::getCurFusion()->oneVal();
  for (auto pt : kParallelTypeTIDs) {
    num_threads = SimplifyingIrBuilder::mulExpr(num_threads, getRaw(pt));
  }
  return num_threads;
}

Expected Kernel Output

Confirm that the expected kernel output strings in the tests accurately reflect the changes made to the kernel launch bounds.

  const std::string expected_kernel = R"(
__global__ void __launch_bounds__(/*MAX_THREADS_PER_BLOCK=*/256) CUDAGeneratedKernel(Tensor<float, 2, 2> T0, Tensor<float, 2, 2> T2) {

jacobhinkle · 2025-01-31T13:45:08Z

!test --diff

jacobhinkle · 2025-02-04T14:48:54Z

!test --diff

jacobhinkle · 2025-02-04T17:16:57Z

Codediff doesn't show anything concerning IMO. Register usage is usually reduced, though sometimes increased by a few registers.

rdspring1 · 2025-02-10T20:02:52Z

csrc/codegen.cpp

    if (kernel_->hasManaged("enable_register_sharing") &&
        kernel_->getManaged<bool>("enable_register_sharing")) {
      NVF_ERROR(
          num_threads_per_cta.has_value(),
          "__launch_bounds__ must be set for register sharing warp specialization");
-      code_ << "__launch_bounds__(/*MAX_THREADS_PER_BLOCK=*/"


Isn't this still necessary because we can determine the number of threads at runtime?

How about changing the predicate to

// Always set __launch_bounds__ when register sharing is enabled. if (kernel_->hasManaged("enable_register_sharing") && kernel_->getManaged<bool>("enable_register_sharing") { ... } else { // Enable __launch_bounds__ when number of threads is known at compile-time. // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast) FusionGuard fg(const_cast<kir::Kernel*>(kernel_)); Val* num_threads = kernel_->summary().parallel_dimension_map.getNumThreadsEachBlock(); if (num_threads->isConstInt()) { code_ << "__launch_bounds__(/*MAX_THREADS_PER_BLOCK=*/" << num_threads->evaluate().as<int64_t>() << ") "; } }

Good call. I now just try and set num_threads_per_cta if it is unset, then I use that to set the launch bounds argument. So in the warp specialization case we will skip the check and just use the provided number of threads.

jacobhinkle · 2025-02-11T14:05:00Z

!test

jacobhinkle · 2025-02-11T15:32:53Z

We hit an error because of this

Fuser/csrc/runtime/executor.cpp

Lines 283 to 286 in fc719a7

    
           // Now that we have launch parameters we can compile the kernel. It's a bit 
        
           // odd we need launch parameters for compilation, need to go back and check 
        
           // why this is the case. 
        
           compiled_kernel_->compile(launch_params.nThreads());

When we compile the kernel, we are launching it one way, but then when we reuse the kernel we can use different launch params. I was thinking that for warp specialization in matmul this is probably not needed because we will have fixed block size. @rdspring1 what is the case where we might need to have runtime determination of block size?

If needed, we could consider caching and recompiling after lowering, as part of CompiledKernel. This is a very similar challenge to index type: see #3850.

rdspring1 · 2025-02-11T16:48:00Z

csrc/codegen.cpp

    if (kernel_->hasManaged("enable_register_sharing") &&
        kernel_->getManaged<bool>("enable_register_sharing")) {
      NVF_ERROR(
          num_threads_per_cta.has_value(),
          "__launch_bounds__ must be set for register sharing warp specialization");
+    }
+    if (num_threads_per_cta.has_value()) {


num_threads_per_cta is always known at compile-time, so the setting launch bounds with a runtime value should be guarded by enable_register_sharing.

Yeah, I get that now. But due to kernel re-use this is not safe. We might have different LaunchParams during compilation than we do for a later launch.

rdspring1 · 2025-02-11T16:48:25Z

The persistent schedulers can use dynamic block sizes. There was interest in exploring warp specialization with register sharing for those schedulers.

https://github.com/NVIDIA/Fuser/blob/main/csrc/scheduler/normalization_inner_outer.cpp#L903-L908

jacobhinkle requested review from zasdfgbnm and rdspring1 January 29, 2025 18:18

zasdfgbnm reviewed Jan 29, 2025

View reviewed changes

rdspring1 reviewed Jan 29, 2025

View reviewed changes

Compute num threads per block properly

eb80538

Fix error caused by using evaluate() without checking if const

b78ae42

Fix failing tests with expected cuda code

254a8a9

jacobhinkle added 2 commits January 31, 2025 08:21

Just use a const_cast

4495654

Disable clang-tidy warning on const_cast

508d7ae

jacobhinkle added 2 commits February 4, 2025 06:58

Merge branch 'main' into jh/set_launch_bounds_more_often

abebb2a

Merge branch 'main' into jh/set_launch_bounds_more_often

1e6e6b9

jacobhinkle requested a review from rdspring1 February 10, 2025 16:39

rdspring1 approved these changes Feb 10, 2025

View reviewed changes

Update logic

0088161

jacobhinkle mentioned this pull request Feb 11, 2025

Add register sharing to warp-specialized circular buffering #3669

Merged

rdspring1 reviewed Feb 11, 2025

View reviewed changes

Set __launch_bounds__ in kernel whenever we are able #3794

Are you sure you want to change the base?

Set __launch_bounds__ in kernel whenever we are able #3794

Uh oh!

Conversation

jacobhinkle commented Jan 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jacobhinkle commented Jan 29, 2025

Uh oh!

github-actions bot commented Jan 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Reviewer Guide 🔍

(Review updated until commit b78ae42)

Uh oh!

zasdfgbnm left a comment

Choose a reason for hiding this comment

Uh oh!

rdspring1 commented Jan 29, 2025

Uh oh!

rdspring1 left a comment

Choose a reason for hiding this comment

Uh oh!

jacobhinkle commented Jan 29, 2025

Uh oh!

jacobhinkle commented Jan 30, 2025

Uh oh!

jacobhinkle commented Jan 31, 2025

Uh oh!

jacobhinkle commented Jan 31, 2025

Uh oh!

github-actions bot commented Jan 31, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Changes walkthrough 📝

PR Reviewer Guide 🔍

Uh oh!

jacobhinkle commented Jan 31, 2025

Uh oh!

jacobhinkle commented Feb 4, 2025

Uh oh!

jacobhinkle commented Feb 4, 2025

Uh oh!

rdspring1 Feb 10, 2025

Choose a reason for hiding this comment

Uh oh!

jacobhinkle Feb 11, 2025

Choose a reason for hiding this comment

Uh oh!

jacobhinkle commented Feb 11, 2025

Uh oh!

jacobhinkle commented Feb 11, 2025

Uh oh!

rdspring1 Feb 11, 2025

Choose a reason for hiding this comment

Uh oh!

jacobhinkle Feb 11, 2025

Choose a reason for hiding this comment

Uh oh!

rdspring1 commented Feb 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

jacobhinkle commented Jan 29, 2025 •

edited

Loading

github-actions bot commented Jan 29, 2025 •

edited

Loading

(Review updated until commit `b78ae42`)

github-actions bot commented Jan 31, 2025 •

edited

Loading

rdspring1 commented Feb 11, 2025 •

edited

Loading