[MLIR] Fix FloorDivSIOpConverter that was failing for index type after the arithmetic op refactor

[Holman]

Summary

ConstantOp should be used instead of ConstantIntOp to be able to support index type.

[Mogball]

Thanks

[Holman]

Can someone help me get this checked in?

mlir/lib/Dialect/Arithmetic/Transforms/ExpandOps.cpp

struct FloorDivSIOpConverter : public OpRewritePattern<arith::FloorDivSIOp> {
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(arith::FloorDivSIOp op,
                                 PatternRewriter &rewriter) const final {
     Location loc = op.getLoc();
     arith::FloorDivSIOp signedFloorDivIOp = cast<arith::FloorDivSIOp>(op);
     Type type = signedFloorDivIOp.getType();
     Value a = signedFloorDivIOp.lhs();
     Value b = signedFloorDivIOp.rhs();
-    Value plusOne = rewriter.create<arith::ConstantIntOp>(loc, 1, type);
-    Value zero = rewriter.create<arith::ConstantIntOp>(loc, 0, type);
-    Value minusOne = rewriter.create<arith::ConstantIntOp>(loc, -1, type);
+    Value plusOne = rewriter.create<arith::ConstantOp>(
+        loc, rewriter.getIntegerAttr(type, 1));
+    Value zero = rewriter.create<arith::ConstantOp>(
+        loc, rewriter.getIntegerAttr(type, 0));
+    Value minusOne = rewriter.create<arith::ConstantOp>(
+        loc, rewriter.getIntegerAttr(type, -1));
     // Compute x = (b<0) ? 1 : -1.
     Value compare =
         rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, b, zero);
     Value x = rewriter.create<SelectOp>(loc, compare, plusOne, minusOne);
     // Compute negative res: -1 - ((x-a)/b).
     Value xMinusA = rewriter.create<arith::SubIOp>(loc, x, a);
     Value xMinusADivB = rewriter.create<arith::DivSIOp>(loc, xMinusA, b);
     Value negRes = rewriter.create<arith::SubIOp>(loc, minusOne, xMinusADivB);

mlir/test/Dialect/Arithmetic/expand-ops.mlir

 // CHECK:           [[TERM1:%.+]] = arith.andi [[NNEG]], [[MNEG]] : i1
 // CHECK:           [[TERM2:%.+]] = arith.andi [[NPOS]], [[MPOS]] : i1
 // CHECK:           [[CMP2:%.+]] = arith.ori [[TERM1]], [[TERM2]] : i1
 // CHECK:           [[RES:%.+]] = select [[CMP2]], [[TRUE3]], [[FALSE3]] : i32
 }
 
 // -----
 
+// Test ceil divide with index type
+// CHECK-LABEL:       func @ceildivi_index
+// CHECK-SAME:     ([[ARG0:%.+]]: index, [[ARG1:%.+]]: index) -> index {
+func @ceildivi_index(%arg0: index, %arg1: index) -> (index) {
+  %res = arith.ceildivsi %arg0, %arg1 : index
+  return %res : index
+
+// CHECK:           [[ONE:%.+]] = arith.constant 1 : index
+// CHECK:           [[ZERO:%.+]] = arith.constant 0 : index
+// CHECK:           [[MINONE:%.+]] = arith.constant -1 : index
+// CHECK:           [[CMP1:%.+]] = arith.cmpi sgt, [[ARG1]], [[ZERO]] : index
+// CHECK:           [[X:%.+]] = select [[CMP1]], [[MINONE]], [[ONE]] : index
+// CHECK:           [[TRUE1:%.+]] = arith.addi [[X]], [[ARG0]] : index
+// CHECK:           [[TRUE2:%.+]] = arith.divsi [[TRUE1]], [[ARG1]] : index
+// CHECK:           [[TRUE3:%.+]] = arith.addi [[ONE]], [[TRUE2]] : index
+// CHECK:           [[FALSE1:%.+]] = arith.subi [[ZERO]], [[ARG0]] : index
+// CHECK:           [[FALSE2:%.+]] = arith.divsi [[FALSE1]], [[ARG1]] : index
+// CHECK:           [[FALSE3:%.+]] = arith.subi [[ZERO]], [[FALSE2]] : index
+// CHECK:           [[NNEG:%.+]] = arith.cmpi slt, [[ARG0]], [[ZERO]] : index
+// CHECK:           [[NPOS:%.+]] = arith.cmpi sgt, [[ARG0]], [[ZERO]] : index
+// CHECK:           [[MNEG:%.+]] = arith.cmpi slt, [[ARG1]], [[ZERO]] : index
+// CHECK:           [[MPOS:%.+]] = arith.cmpi sgt, [[ARG1]], [[ZERO]] : index
+// CHECK:           [[TERM1:%.+]] = arith.andi [[NNEG]], [[MNEG]] : i1
+// CHECK:           [[TERM2:%.+]] = arith.andi [[NPOS]], [[MPOS]] : i1
+// CHECK:           [[CMP2:%.+]] = arith.ori [[TERM1]], [[TERM2]] : i1
+// CHECK:           [[RES:%.+]] = select [[CMP2]], [[TRUE3]], [[FALSE3]] : index
+}
+
+// -----
+
 // Test floor divide with signed integer
 // CHECK-LABEL:       func @floordivi
 // CHECK-SAME:     ([[ARG0:%.+]]: i32, [[ARG1:%.+]]: i32) -> i32 {
 func @floordivi(%arg0: i32, %arg1: i32) -> (i32) {
   %res = arith.floordivsi %arg0, %arg1 : i32
   return %res : i32
 // CHECK:           [[ONE:%.+]] = arith.constant 1 : i32
 // CHECK:           [[ZERO:%.+]] = arith.constant 0 : i32
 // CHECK:           [[MIN1:%.+]] = arith.constant -1 : i32
 // CHECK:           [[CMP1:%.+]] = arith.cmpi slt, [[ARG1]], [[ZERO]] : i32
 // CHECK:           [[X:%.+]] = select [[CMP1]], [[ONE]], [[MIN1]] : i32
 // CHECK:           [[TRUE1:%.+]] = arith.subi [[X]], [[ARG0]] : i32
 // CHECK:           [[TRUE2:%.+]] = arith.divsi [[TRUE1]], [[ARG1]] : i32
 // CHECK:           [[TRUE3:%.+]] = arith.subi [[MIN1]], [[TRUE2]] : i32
 // CHECK:           [[FALSE:%.+]] = arith.divsi [[ARG0]], [[ARG1]] : i32
 // CHECK:           [[NNEG:%.+]] = arith.cmpi slt, [[ARG0]], [[ZERO]] : i32
 // CHECK:           [[NPOS:%.+]] = arith.cmpi sgt, [[ARG0]], [[ZERO]] : i32
 // CHECK:           [[MNEG:%.+]] = arith.cmpi slt, [[ARG1]], [[ZERO]] : i32
 // CHECK:           [[MPOS:%.+]] = arith.cmpi sgt, [[ARG1]], [[ZERO]] : i32
 // CHECK:           [[TERM1:%.+]] = arith.andi [[NNEG]], [[MPOS]] : i1
 // CHECK:           [[TERM2:%.+]] = arith.andi [[NPOS]], [[MNEG]] : i1
 // CHECK:           [[CMP2:%.+]] = arith.ori [[TERM1]], [[TERM2]] : i1
 // CHECK:           [[RES:%.+]] = select [[CMP2]], [[TRUE3]], [[FALSE]] : i32
 }
+
+// -----
+
+// Test floor divide with index type
+// CHECK-LABEL:       func @floordivi_index
+// CHECK-SAME:     ([[ARG0:%.+]]: index, [[ARG1:%.+]]: index) -> index {
+func @floordivi_index(%arg0: index, %arg1: index) -> (index) {
+  %res = arith.floordivsi %arg0, %arg1 : index
+  return %res : index
+// CHECK:           [[ONE:%.+]] = arith.constant 1 : index
+// CHECK:           [[ZERO:%.+]] = arith.constant 0 : index
+// CHECK:           [[MIN1:%.+]] = arith.constant -1 : index
+// CHECK:           [[CMP1:%.+]] = arith.cmpi slt, [[ARG1]], [[ZERO]] : index
+// CHECK:           [[X:%.+]] = select [[CMP1]], [[ONE]], [[MIN1]] : index
+// CHECK:           [[TRUE1:%.+]] = arith.subi [[X]], [[ARG0]] : index
+// CHECK:           [[TRUE2:%.+]] = arith.divsi [[TRUE1]], [[ARG1]] : index
+// CHECK:           [[TRUE3:%.+]] = arith.subi [[MIN1]], [[TRUE2]] : index
+// CHECK:           [[FALSE:%.+]] = arith.divsi [[ARG0]], [[ARG1]] : index
+// CHECK:           [[NNEG:%.+]] = arith.cmpi slt, [[ARG0]], [[ZERO]] : index
+// CHECK:           [[NPOS:%.+]] = arith.cmpi sgt, [[ARG0]], [[ZERO]] : index
+// CHECK:           [[MNEG:%.+]] = arith.cmpi slt, [[ARG1]], [[ZERO]] : index
+// CHECK:           [[MPOS:%.+]] = arith.cmpi sgt, [[ARG1]], [[ZERO]] : index
+// CHECK:           [[TERM1:%.+]] = arith.andi [[NNEG]], [[MPOS]] : i1
+// CHECK:           [[TERM2:%.+]] = arith.andi [[NPOS]], [[MNEG]] : i1
+// CHECK:           [[CMP2:%.+]] = arith.ori [[TERM1]], [[TERM2]] : i1
+// CHECK:           [[RES:%.+]] = select [[CMP2]], [[TRUE3]], [[FALSE]] : index
+}