diff --git a/mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp b/mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp --- a/mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp +++ b/mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp @@ -653,8 +653,8 @@ // affine.apply here so that it composes nicely with the provided map. AffineMap stepMap = AffineMap::get(0, 3, - (rewriter.getAffineSymbolExpr(0) - - rewriter.getAffineSymbolExpr(1).ceilDiv( + ((rewriter.getAffineSymbolExpr(0) - + rewriter.getAffineSymbolExpr(1)).ceilDiv( rewriter.getAffineSymbolExpr(2)))); Value launchBound = rewriter.create( loc, annotation.boundMap.compose(stepMap), diff --git a/mlir/test/Conversion/LoopsToGPU/parallel_loop.mlir b/mlir/test/Conversion/LoopsToGPU/parallel_loop.mlir --- a/mlir/test/Conversion/LoopsToGPU/parallel_loop.mlir +++ b/mlir/test/Conversion/LoopsToGPU/parallel_loop.mlir @@ -15,7 +15,7 @@ return } -// CHECK: #[[MAP0:.*]] = affine_map<()[s0, s1, s2] -> (s0 - s1 ceildiv s2)> +// CHECK: #[[MAP0:.*]] = affine_map<()[s0, s1, s2] -> ((s0 - s1) ceildiv s2)> // CHECK: #[[MAP1:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s0 + s1)> // CHECK: module { @@ -66,7 +66,7 @@ return } -// CHECK: #[[MAP0:.*]] = affine_map<()[s0, s1, s2] -> (s0 - s1 ceildiv s2)> +// CHECK: #[[MAP0:.*]] = affine_map<()[s0, s1, s2] -> ((s0 - s1) ceildiv s2)> // CHECK: #[[MAP1:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s0 + s1)> // CHECK: module { @@ -115,7 +115,7 @@ return } -// CHECK: #[[MAP0:.*]] = affine_map<()[s0, s1, s2] -> (s0 - s1 ceildiv s2)> +// CHECK: #[[MAP0:.*]] = affine_map<()[s0, s1, s2] -> ((s0 - s1) ceildiv s2)> // CHECK: #[[MAP1:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s0 + s1)> // CHECK: module { @@ -166,7 +166,7 @@ return } -// CHECK: #[[MAP0:.*]] = affine_map<()[s0, s1, s2] -> (s0 - s1 ceildiv s2)> +// CHECK: #[[MAP0:.*]] = affine_map<()[s0, s1, s2] -> ((s0 - s1) ceildiv s2)> // CHECK: #[[MAP1:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s0 + s1)> // CHECK: module { @@ -241,7 +241,7 @@ } // CHECK: #[[MAP0:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)> -// CHECK: #[[MAP1:.*]] = affine_map<()[s0, s1, s2] -> (s0 - s1 ceildiv s2)> +// CHECK: #[[MAP1:.*]] = affine_map<()[s0, s1, s2] -> ((s0 - s1) ceildiv s2)> // CHECK: #[[MAP2:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s0 + s1)> // CHECK: #[[MAP3:.*]] = affine_map<(d0)[s0] -> (2, -d0 + s0)> // CHECK: #[[MAP4:.*]] = affine_map<(d0)[s0] -> (3, -d0 + s0)>