Commit 1940b3c9 authored by Lei Wang's avatar Lei Wang Committed by LeiWang1999
Browse files

[Bugfix] Fix a bug when simplifying warp combination for T.gemm (#540)

parent a65f481e
...@@ -194,10 +194,6 @@ std::pair<int, int> Gemm::ComputeWarpPartition(int num_warps, Target target, ...@@ -194,10 +194,6 @@ std::pair<int, int> Gemm::ComputeWarpPartition(int num_warps, Target target,
// Try all possible combinations that satisfy the constraints // Try all possible combinations that satisfy the constraints
for (int m = 1; m <= max_m_warps && m <= num_warps; m++) { for (int m = 1; m <= max_m_warps && m <= num_warps; m++) {
int n = num_warps / m; int n = num_warps / m;
if (n > max_n_warps)
continue;
if (m * n != num_warps)
continue;
// Calculate how balanced this partition is // Calculate how balanced this partition is
float m_per_warp = static_cast<float>(this->M) / (m * kMPerWarp); float m_per_warp = static_cast<float>(this->M) / (m * kMPerWarp);
......
...@@ -179,6 +179,9 @@ class _JitImplementation: ...@@ -179,6 +179,9 @@ class _JitImplementation:
else: else:
raise ValueError(f"Invalid function type: {type(program_result_source)}") raise ValueError(f"Invalid function type: {type(program_result_source)}")
if self.verbose:
logger.info(f"Verbose: Compiling for program \n {program_result.script()}")
kernel_result = compile( kernel_result = compile(
program_result, program_result,
out_idx=self.out_idx, out_idx=self.out_idx,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment