Skip to content

Commit

Permalink
Ensure expression simplification is enabled in proveLinearAndGetStride
Browse files Browse the repository at this point in the history
When debugging the matmul scheduler, I noticed that we hit an error when
generating Hopper MMA when we use `NVFUSER_DISABLE=expr_simplify`. This
PR makes that possible again. Note that the option is disabled only
within `proveLinearAndGetStride` so it has no direct effect on index
expressions, which is what we want to look at when using that option.
  • Loading branch information
jacobhinkle committed Nov 5, 2024
1 parent 162a13b commit bc8aabd
Showing 1 changed file with 5 additions and 0 deletions.
5 changes: 5 additions & 0 deletions csrc/device_lower/utils.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1907,6 +1907,11 @@ Val* proveLinearAndGetStride(
const ValGroup& linear_g,
const ValGroups& domain) {
FusionGuard fg(linear_g->front()->fusion());
// This function uses simplifyExpr extensively. If we have disable expression
// simplification in order to help inspect generated kernels then we will get
// incorrect results here. Instead, we ensure it is enabled using this guard.
DisableOptionsGuard dog;
DisableOptionsGuard::getCurOptions().unset(DisableOption::ExprSimplify);
if (simplifyExpr(extent(linear_g))->isOne()) {
// If the extent of the linear group is 1, we always consider it as linear,
// regardless of its relationship with domain. For this case, we use stride
Expand Down

0 comments on commit bc8aabd

Please sign in to comment.