diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -679,8 +679,8 @@ return; // Simplify the child expressions first. - auto lhs = binExpr.getLHS(); - auto rhs = binExpr.getRHS(); + AffineExpr lhs = binExpr.getLHS(); + AffineExpr rhs = binExpr.getRHS(); simplifyExprAndOperands(lhs, operands); simplifyExprAndOperands(rhs, operands); expr = getAffineBinaryOpExpr(binExpr.getKind(), lhs, rhs); @@ -691,11 +691,18 @@ return; } + // The `lhs` and `rhs` may be different post construction of simplified expr. + lhs = binExpr.getLHS(); + rhs = binExpr.getRHS(); auto rhsConst = rhs.dyn_cast(); if (!rhsConst) return; int64_t rhsConstVal = rhsConst.getValue(); + // Undefined exprsessions aren't touched; IR can still be valid with them. + if (rhsConstVal == 0) + return; + AffineExpr quotientTimesDiv, rem; int64_t divisor; diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir --- a/mlir/test/Dialect/Affine/canonicalize.mlir +++ b/mlir/test/Dialect/Affine/canonicalize.mlir @@ -1155,6 +1155,8 @@ // Simplification of maps exploiting operand info. +// CHECK: #[[$MAP_SIMPLER:.*]] = affine_map<(d0, d1) -> (((d0 + d1) mod 458313) floordiv 227)> + // CHECK-LABEL: func @simplify_with_operands func.func @simplify_with_operands(%N: index, %A: memref) { // CHECK-NEXT: affine.for %[[I:.*]] = 0 to %{{.*}} @@ -1186,5 +1188,15 @@ "test.foo"(%x) : (f32) -> () } + affine.for %arg0 = 0 to %N step 128 { + affine.for %arg4 = 0 to 32 step 32 { + affine.for %arg5 = 0 to 128 { + // CHECK: affine.apply #[[$MAP_SIMPLER]] + %x = affine.apply affine_map<(d0, d1, d2) -> (((d0 + d2) mod 458313) floordiv 227 + d1 floordiv 256)>(%arg0, %arg4, %arg5) + "test.foo"(%x) : (index) -> () + } + } + } + return }