feat(meta_schedule): expand CUDA unroll steps for SM70 optimization

[jianhua1724]

Motivation

Expand CUDA unroll search space to support optimal loop unrolling sizes for SM70 (V100) GPUs.
Add missing critical unroll steps: 32, 128, 256.

Changes

Only modify 1 line:
Extend unroll_max_steps in ScheduleRule::DefaultCUDA()
From: {0, 16, 64, 512, 1024}
To: {0, 16, 32, 64, 128, 256, 512, 1024}

Benefits

• 5%~15% performance improvement on SM70 (V100)
• No breaking changes
• Backward compatible with all CUDA architectures
• Only expand search space, no logic modification

Test

Compiled and tested on SM70 (V100), no regression.

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances the MetaSchedule's CUDA optimization capabilities by refining the loop unrolling strategy. By introducing a broader range of unroll steps, the system can now discover more efficient execution paths, particularly benefiting SM70 (V100) GPUs. This targeted adjustment aims to boost performance without impacting compatibility or stability across different CUDA architectures.

Highlights

• CUDA Unroll Search Space Expansion: Expanded the unroll_max_steps in ScheduleRule::DefaultCUDA() to include additional critical unroll steps (32, 128, 256) for SM70 (V100) GPUs.
• Performance Improvement: This change is expected to deliver a 5% to 15% performance improvement on SM70 (V100) GPUs by allowing for more optimal loop unrolling sizes.
• Compatibility and Safety: The modification is backward compatible with all CUDA architectures, introduces no breaking changes, and only expands the search space without altering existing logic.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request expands the unroll_max_steps array in the DefaultCUDA schedule rule for ParallelizeVectorizeUnroll in src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc, adding intermediate unroll steps (32, 128, 256) to the search space. The feedback suggests that while this expansion is beneficial for newer CUDA architectures (SM70+), applying it universally could increase auto-tuning time for older architectures. It is recommended to make the unroll_max_steps conditional on the target's compute capability to optimize tuning time.

[gemini-code-assist]

This change expands the unroll search space for all CUDA targets, which can increase auto-tuning time for architectures where these new unroll steps are not beneficial.

To make this optimization more targeted to SM70+ GPUs as intended, I suggest making the choice of unroll steps conditional on the target's compute capability. This avoids slowing down tuning on older architectures. This can be done cleanly using an immediately-invoked lambda expression.

The condition can be sm_version >= 70 if the optimization is beneficial for SM70 and newer, or sm_version == 70 if it's specific to V100.

          /*unroll_max_steps=*/[]() -> ffi::Array<Integer> {
            auto target = tvm::Target::Current(true);
            if (target.defined() && target->kind->name == "cuda") {
              if (const auto* sm_ptr = target->GetAttr<Integer>("sm")) {
                if (sm_ptr->value() >= 70) {
                  return {0, 16, 32, 64, 128, 256, 512, 1024};
                }
              }
            }
            return {0, 16, 64, 512, 1024};
          }(),

[tlopex]

Thanks for the contribution. The motivation makes sense, but we need concrete data before merging a change to the default search space that affects all CUDA users:

1. Which model/workload did you benchmark? Please share actual latency numbers (before vs after) on V100.
2. The search space grows from 5 to 8 candidates. Did you measure the tuning time increase?
3. "No regression" on other architectures — does this mean you tested on non-SM70 GPUs, or just that the code compiles?

[jianhua1724]

Thanks for the contribution. The motivation makes sense, but we need concrete data before merging a change to the default search space that affects all CUDA users:

1. Which model/workload did you benchmark? Please share actual latency numbers (before vs after) on V100.
2. The search space grows from 5 to 8 candidates. Did you measure the tuning time increase?
3. "No regression" on other architectures — does this mean you tested on non-SM70 GPUs, or just that the code compiles?

Thank you for your reply. I will answer your three questions as follows:
My scenario is to accelerate the model training of algorithm engineers on the V100 GPU (and inference acceleration will be supported later).
The computationally heavy network structure in the model is a feed-forward network with fixed input shapes, such as the following:

batch_size = 1024
input_size = 256
hidden_size = 512
output_size = 256
class DeepFFN(nn.Module):
def init(self):
super().init()
self.fc1 = nn.Linear(input_size, hidden_size, dtype=torch.float16)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_size, output_size, dtype=torch.float16)

def forward(self, x):
    return self.fc2(self.relu(self.fc1(x)))

I defined three operators to replace PyTorch's native implementation for training acceleration.
The three operators are: one for the forward propagation of the network, and two for the backward propagation logic (the backward propagation is complex, so I split it into two operators).
Then I used TVM's MetaSchedule to compile the three operators and finally replaced PyTorch's native implementation.

1.Test data:
1.1 Speed difference of a single operator:
The following is the test data of the forward propagation operator in the two versions:
Version with "unroll_max_steps=ffi::Array{0, 16, 64, 512, 1024}":

2026-04-13 12:33:33 [INFO] [task_scheduler.cc:326]
ID | Name | FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done

0 | main | 537657344 | 1 | 14873.9572 | 36.1476 | 36.1476 | 100 | Y

Total trials: 100
Total latency (us): 36.1476

Version with "unroll_max_steps=ffi::Array{0, 16, 32, 64, 128, 256, 512, 1024}":

ID | Name | FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done

0 | main | 537657344 | 1 | 17980.4812 | 29.9023 | 29.9023 | 100 | Y

Total trials: 100
Total latency (us): 29.9023

It can be seen that the latency is reduced by approximately 17.27%.

1.2 Speed difference in total training:
Training time for num_epochs = 1000:
PyTorch time cost: 1.1044s
Training time cost after TVM operator acceleration: 0.9143s (unroll_max_steps=ffi::Array{0, 16, 64, 512, 1024} version)
Training time cost after TVM operator acceleration: 0.7945s (unroll_max_steps=ffi::Array{0, 16, 32, 64, 128, 256, 512, 1024} version)

It can be seen that the unroll_max_steps=ffi::Array{0, 16, 32, 64, 128, 256, 512, 1024} version is approximately 13% faster than the previous version.

Other parameters in the test are as follows:

target = Target(
"cuda -arch=sm_70 "
"-max_num_threads=1024 "
"-max_threads_per_block=1024 "
"-thread_warp_size=32 "
"-max_shared_memory_per_block=98304 "
"-registers_per_block=65536"
"+use_tensor_core"
)
tune_tir(
mod=mod,
target=target,
database=database,
max_trials_global=100,
num_trials_per_iter=8,
space="cuda-tensorcore",
work_dir=work_dir,
)
num_epochs = 1000

2.According to the test results, this modification does not increase the tuning time.
For three operators with max_trials_global=100 for each operator, the total tuning time is:
Version with "unroll_max_steps=ffi::Array{0, 16, 64, 512, 1024}": 67 minutes 7 seconds.
Version with "unroll_max_steps=ffi::Array{0, 16, 32, 64, 128, 256, 512, 1024}": 63 minutes 43 seconds.
There is no significant time difference.

From the source code analysis:It only increases the mutation selection possibilities of MutateUnrollNode in mutate_unroll.cc (i.e., it affects the diversity of the search space).
The parameter that truly affects TVM tuning time is max_trials_global in:

tune_tir(
mod=mod,
target=target,
database=database,
max_trials_global=max_trials_global,
num_trials_per_iter=num_trials_per_iter,
space="cuda-tensorcore",
work_dir=work_dir,
)
It is not the size of the unroll_max_steps array.

3.Currently, I only have a V100 GPU. There are many other CUDA GPUs, and it is difficult for me to obtain all of them for comprehensive comparison and testing due to the high financial cost.
This is also a practical problem that most community contributors face.
However, I have carefully and thoroughly analyzed the MetaSchedule source code.
I have analyzed the complete processing link from ParallelizeVectorizeUnrollNode::Apply to MutateUnrollNode::FindUnrollDecision.
The unroll_max_steps parameter is used in these two key places: the former sets candidate values during initial scheduling, and the latter performs mutation selection during tuning.
This processing mechanism is unified for all CUDA GPUs. The parameter itself only provides the search space and does not contain hardware-specific logic, so it has excellent cross-architecture compatibility.

[tlopex]

Thanks for providing that! Approved!