[Mlir-commits] [mlir] d17428d - [MLIR][GPU] fix loop trip count computation in LoopsToGPU

[Stephan Herhut]

Author: Stephan Herhut
Date: 2020-03-02T15:53:33+01:00
New Revision: d17428d9513555535268104a94c11ddc6c6c6243

URL: https://github.com/llvm/llvm-project/commit/d17428d9513555535268104a94c11ddc6c6c6243
DIFF: https://github.com/llvm/llvm-project/commit/d17428d9513555535268104a94c11ddc6c6c6243.diff

LOG: [MLIR][GPU] fix loop trip count computation in LoopsToGPU

Summary: Added brackets to fix the loop trip count computation.
The brackets ensure the bounds are subtracted before we divide
the result by the step of the loop.

Differential Revision: https://reviews.llvm.org/D75449

Added: 
    

Modified: 
    mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp
    mlir/test/Conversion/LoopsToGPU/parallel_loop.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp b/mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp
index 5b6b3a2d4f56..ed9400fc2ad0 100644
--- a/mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp
+++ b/mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp
@@ -653,8 +653,8 @@ static LogicalResult processParallelLoop(ParallelOp parallelOp,
           // 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<AffineApplyOp>(
               loc, annotation.boundMap.compose(stepMap),

diff  --git a/mlir/test/Conversion/LoopsToGPU/parallel_loop.mlir b/mlir/test/Conversion/LoopsToGPU/parallel_loop.mlir
index b4b91453a60a..2a440a4456ba 100644
--- a/mlir/test/Conversion/LoopsToGPU/parallel_loop.mlir
+++ b/mlir/test/Conversion/LoopsToGPU/parallel_loop.mlir
@@ -15,7 +15,7 @@ func @parallel_loop_bidy_bidx(%arg0 : index, %arg1 : index, %arg2 : index,
   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 @@ func @parallel_loop_tiled(%arg0 : index, %arg1 : index, %arg2 : index,
   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 @@ func @parallel_loop_bidy_seq(%arg0 : index, %arg1 : index, %arg2 : index,
   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 @@ func @parallel_loop_tiled_seq(%arg0 : index, %arg1 : index, %arg2 : index,
   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 @@ module {
 }
 
 // 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)>