[Mlir-commits] [mlir] [mlir][affine][gpu] Replace DivSIOp to CeilDivSIOp when lowering to GPU launch (PR #73328)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Hsiangkai Wang (Hsiangkai)

<details>
<summary>Changes</summary>

When converting affine.for to GPU launch operator, we have to calculate the block dimension and thread dimension for the launch operator.

The formula of the dimension size is

(upper_bound - lower_bound) / step_size

When the difference is indivisible by step_size, we use rounding-to-zero as the division result. However, the block dimension and thread dimension is right-open range, i.e., [0, block_dim) and [0, thread_dim). So, we will get the wrong result if we use DivSIOp. In this patch, we replace it with CeilDivSIOp to get the correct block and thread dimension values.

---
Full diff: https://github.com/llvm/llvm-project/pull/73328.diff


2 Files Affected:

- (modified) mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp (+1-1) 
- (modified) mlir/test/Conversion/SCFToGPU/step_positive.mlir (+2-2) 


``````````diff
diff --git a/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp b/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp
index 11b4cbb2506705b..a4862ebb24b178b 100644
--- a/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp
+++ b/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp
@@ -195,7 +195,7 @@ AffineLoopToGpuConverter::collectBounds(AffineForOp forOp, unsigned numLoops) {
                                                 upperBound, lowerBound);
     Value step = getOrCreateStep(currentLoop, builder);
     if (getConstantIntValue(step) != static_cast<int64_t>(1))
-      range = builder.create<arith::DivSIOp>(currentLoop.getLoc(), range, step);
+      range = builder.create<arith::CeilDivSIOp>(currentLoop.getLoc(), range, step);
     dims.push_back(range);
 
     lbs.push_back(lowerBound);
diff --git a/mlir/test/Conversion/SCFToGPU/step_positive.mlir b/mlir/test/Conversion/SCFToGPU/step_positive.mlir
index 97fd7d598621b39..84e8454e56171de 100644
--- a/mlir/test/Conversion/SCFToGPU/step_positive.mlir
+++ b/mlir/test/Conversion/SCFToGPU/step_positive.mlir
@@ -3,8 +3,8 @@
 // CHECK-LABEL: @step_var
 func.func @step_var(%A : memref<?x?xf32>, %B : memref<?x?xf32>) {
   // Check that we divide by step.
-  // CHECK:  %[[range_i:.*]] = arith.divsi {{.*}}, %{{.*}}
-  // CHECK:  %[[range_j:.*]] = arith.divsi {{.*}}, %{{.*}}
+  // CHECK:  %[[range_i:.*]] = arith.ceildivsi {{.*}}, %{{.*}}
+  // CHECK:  %[[range_j:.*]] = arith.ceildivsi {{.*}}, %{{.*}}
 
   // CHECK: gpu.launch
   // CHECK-SAME: blocks(%{{[^)]*}}, %{{[^)]*}}, %{{[^)]*}}) in (%{{[^)]*}} = %[[range_i]], %{{[^)]*}} = %{{[^)]*}}, %{{[^)]*}} = %{{[^)]*}})

``````````

</details>


https://github.com/llvm/llvm-project/pull/73328