use arith instead of index

Author: CoTinker

mlir/include/mlir/Conversion/Passes.td

@@ -618,7 +618,6 @@ def ConvertGpuOpsToNVVMOps : Pass<"convert-gpu-to-nvvm", "gpu::GPUModuleOp"> {
     "cf::ControlFlowDialect",
     "memref::MemRefDialect",
     "NVVM::NVVMDialect",
-    "index::IndexDialect"
   ];
   let options = [
     Option<"indexBitwidth", "index-bitwidth", "unsigned",
@@ -648,7 +647,6 @@ def ConvertGpuOpsToROCDLOps : Pass<"convert-gpu-to-rocdl", "gpu::GPUModuleOp"> {
     "amdgpu::AMDGPUDialect",
     "cf::ControlFlowDialect",
     "memref::MemRefDialect",
-    "index::IndexDialect"
   ];
   let options = [
     Option<"chipset", "chipset", "std::string",

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

@@ -23,7 +23,6 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
-#include "mlir/Dialect/Index/IR/IndexDialect.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

@@ -31,7 +31,6 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
-#include "mlir/Dialect/Index/IR/IndexDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"

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

@@ -11,9 +11,9 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
-#include "mlir/Dialect/Index/IR/IndexOps.h"
 #include "mlir/IR/PatternMatch.h"
 
 using namespace mlir;
@@ -26,13 +26,15 @@ struct GpuGlobalIdRewriter : public OpRewritePattern<gpu::GlobalIdOp> {
                                 PatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
     auto dim = op.getDimension();
-    auto blockId = gpu::BlockIdOp::create(rewriter, loc, dim);
-    auto blockDim = gpu::BlockDimOp::create(rewriter, loc, dim);
+    Value blockId = gpu::BlockIdOp::create(rewriter, loc, dim);
+    Value blockDim = gpu::BlockDimOp::create(rewriter, loc, dim);
+    auto indexType = rewriter.getIndexType();
     // Compute blockId.x * blockDim.x
-    auto tmp = index::MulOp::create(rewriter, op.getLoc(), blockId, blockDim);
-    auto threadId = gpu::ThreadIdOp::create(rewriter, loc, dim);
+    Value tmp =
+        arith::MulIOp::create(rewriter, loc, indexType, blockId, blockDim);
+    Value threadId = gpu::ThreadIdOp::create(rewriter, loc, dim);
     // Compute threadId.x + blockId.x * blockDim.x
-    rewriter.replaceOpWithNewOp<index::AddOp>(op, threadId, tmp);
+    rewriter.replaceOpWithNewOp<arith::AddIOp>(op, indexType, threadId, tmp);
     return success();
   }
 };

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

@@ -8,27 +8,27 @@ module {
                threads(%tx, %ty, %tz) in (%block_x = %sz, %block_y = %sz, %block_z = %sz) {
       // CHECK: %[[BIDY:.*]] = gpu.block_id x
       // CHECK-NEXT: %[[BDIMY:.*]] = gpu.block_dim x
-      // CHECK-NEXT: %[[TMPY:.*]] = index.mul %[[BIDY]], %[[BDIMY]]
+      // CHECK-NEXT: %[[TMPY:.*]] = arith.muli %[[BIDY]], %[[BDIMY]]
       // CHECK-NEXT: %[[TIDX:.*]] = gpu.thread_id x
-      // CHECK-NEXT: %[[GIDX:.*]] = index.add %[[TIDX]], %[[TMPY]]
+      // CHECK-NEXT: %[[GIDX:.*]] = arith.addi %[[TIDX]], %[[TMPY]]
       %idx = gpu.global_id x
       // CHECK: memref.store %[[GIDX]], %[[MEM]][] : memref<index, 1>
       memref.store %idx, %mem[] : memref<index, 1>
 
       // CHECK: %[[BIDY:.*]] = gpu.block_id y
       // CHECK-NEXT: %[[BDIMY:.*]] = gpu.block_dim y
-      // CHECK-NEXT: %[[TMPY:.*]] = index.mul %[[BIDY]], %[[BDIMY]]
+      // CHECK-NEXT: %[[TMPY:.*]] = arith.muli %[[BIDY]], %[[BDIMY]]
       // CHECK-NEXT: %[[TIDY:.*]] = gpu.thread_id y
-      // CHECK-NEXT: %[[GIDY:.*]] = index.add %[[TIDY]], %[[TMPY]]
+      // CHECK-NEXT: %[[GIDY:.*]] = arith.addi %[[TIDY]], %[[TMPY]]
       %idy = gpu.global_id y
       // CHECK: memref.store %[[GIDY]], %[[MEM]][] : memref<index, 1>
       memref.store %idy, %mem[] : memref<index, 1>
 
       // CHECK: %[[BIDZ:.*]] = gpu.block_id z
       // CHECK-NEXT: %[[BDIMZ:.*]] = gpu.block_dim z
-      // CHECK-NEXT: %[[TMPZ:.*]] = index.mul %[[BIDZ]], %[[BDIMZ]]
+      // CHECK-NEXT: %[[TMPZ:.*]] = arith.muli %[[BIDZ]], %[[BDIMZ]]
       // CHECK-NEXT: %[[TIDZ:.*]] = gpu.thread_id z
-      // CHECK-NEXT: %[[GIDZ:.*]] = index.add %[[TIDZ]], %[[TMPZ]]
+      // CHECK-NEXT: %[[GIDZ:.*]] = arith.addi %[[TIDZ]], %[[TMPZ]]
       %idz = gpu.global_id z
       // CHECK: memref.store %[[GIDZ]], %[[MEM]][] : memref<index, 1>
       memref.store %idz, %mem[] : memref<index, 1>