[MLIR][GPU] Use arith instead of index for subgroup_id (#137843)

Trying to simplify situation by using `arith` dialect instead of `index`
in the rewriting of `gpu.subgroup_id`.

Author: lialan

mlir/lib/Dialect/GPU/Transforms/SubgroupIdRewriter.cpp

@@ -53,23 +53,25 @@ struct GpuSubgroupIdRewriter final : OpRewritePattern<gpu::SubgroupIdOp> {
     //             subgroup_size
 
     Location loc = op->getLoc();
+    Type indexType = rewriter.getIndexType();
 
     Value dimX = rewriter.create<gpu::BlockDimOp>(loc, gpu::Dimension::x);
     Value dimY = rewriter.create<gpu::BlockDimOp>(loc, gpu::Dimension::y);
     Value tidX = rewriter.create<gpu::ThreadIdOp>(loc, gpu::Dimension::x);
     Value tidY = rewriter.create<gpu::ThreadIdOp>(loc, gpu::Dimension::y);
     Value tidZ = rewriter.create<gpu::ThreadIdOp>(loc, gpu::Dimension::z);
 
-    Value dimYxIdZ = rewriter.create<index::MulOp>(loc, dimY, tidZ);
-    Value dimYxIdZPlusIdY = rewriter.create<index::AddOp>(loc, dimYxIdZ, tidY);
+    Value dimYxIdZ = rewriter.create<arith::MulIOp>(loc, indexType, dimY, tidZ);
+    Value dimYxIdZPlusIdY =
+        rewriter.create<arith::AddIOp>(loc, indexType, dimYxIdZ, tidY);
     Value dimYxIdZPlusIdYTimesDimX =
-        rewriter.create<index::MulOp>(loc, dimX, dimYxIdZPlusIdY);
-    Value IdXPlusDimYxIdZPlusIdYTimesDimX =
-        rewriter.create<index::AddOp>(loc, tidX, dimYxIdZPlusIdYTimesDimX);
+        rewriter.create<arith::MulIOp>(loc, indexType, dimX, dimYxIdZPlusIdY);
+    Value IdXPlusDimYxIdZPlusIdYTimesDimX = rewriter.create<arith::AddIOp>(
+        loc, indexType, tidX, dimYxIdZPlusIdYTimesDimX);
     Value subgroupSize = rewriter.create<gpu::SubgroupSizeOp>(
         loc, rewriter.getIndexType(), /*upper_bound = */ nullptr);
-    Value subgroupIdOp = rewriter.create<index::DivUOp>(
-        loc, IdXPlusDimYxIdZPlusIdYTimesDimX, subgroupSize);
+    Value subgroupIdOp = rewriter.create<arith::DivUIOp>(
+        loc, indexType, IdXPlusDimYxIdZPlusIdYTimesDimX, subgroupSize);
     rewriter.replaceOp(op, {subgroupIdOp});
     return success();
   }

mlir/test/Dialect/GPU/subgroupId-rewrite.mlir

@@ -10,12 +10,12 @@ func.func @subgroupId(%sz : index, %mem: memref<index, 1>) {
     // CHECK-NEXT: %[[TIDX:.*]] = gpu.thread_id  x
     // CHECK-NEXT: %[[TIDY:.*]] = gpu.thread_id  y
     // CHECK-NEXT: %[[TIDZ:.*]] = gpu.thread_id  z
-    // CHECK-NEXT: %[[T0:.*]] = index.mul %[[DIMY]], %[[TIDZ]]
-    // CHECK-NEXT: %[[T1:.*]] = index.add %[[T0]], %[[TIDY]]
-    // CHECK-NEXT: %[[T2:.*]] = index.mul %[[DIMX]], %[[T1]]
-    // CHECK-NEXT: %[[T3:.*]] = index.add %[[TIDX]], %[[T2]]
+    // CHECK-NEXT: %[[T0:.*]] = arith.muli %[[DIMY]], %[[TIDZ]] : index
+    // CHECK-NEXT: %[[T1:.*]] = arith.addi %[[T0]], %[[TIDY]] : index
+    // CHECK-NEXT: %[[T2:.*]] = arith.muli %[[DIMX]], %[[T1]] : index
+    // CHECK-NEXT: %[[T3:.*]] = arith.addi %[[TIDX]], %[[T2]] : index
     // CHECK-NEXT: %[[T4:.*]] = gpu.subgroup_size : index
-    // CHECK-NEXT: %[[T5:.*]] = index.divu %[[T3]], %[[T4]]
+    // CHECK-NEXT: %[[T5:.*]] = arith.divui %[[T3]], %[[T4]] : index
     %idz = gpu.subgroup_id : index
     memref.store %idz, %mem[] : memref<index, 1>
     gpu.terminator