# Folder for division by one.

[phoenix-meadowlark]

Summary

• Adds a folder for integer division by one with the divi_signed and divi_unsigned ops.
• Creates tests for scalar and tensor versions of these ops.
• Modifies the test in parallel-loop-collapsing.mlir so that it doesn't assume division by one will be in the output.

[rriddle]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1725

This comment looks weirdly placed, can you move it to a more appropriate location?

2261

Same here.

[phoenix-meadowlark]

Moved comments to a better location.

[mehdi_amini]

Do you need someone to commit this for you? (if so please mention it in the comment in general)

[phoenix-meadowlark]

In D78518#2006079, @mehdi_amini wrote:

Do you need someone to commit this for you? (if so please mention it in the comment in general)

I do! Thank you for letting me know. I will in the future.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

OpFoldResult SignedDivIOp::fold(ArrayRef<Attribute> operands) {
   bool overflowOrDiv0 = false;
   auto result = constFoldBinaryOp<IntegerAttr>(operands, [&](APInt a, APInt b) {
     if (overflowOrDiv0 || !b) {
       overflowOrDiv0 = true;
       return a;
     }
     return a.sdiv_ov(b, overflowOrDiv0);
   });
+
+  // Fold out division by one.
+  if (auto rhs = operands[1].dyn_cast_or_null<IntegerAttr>()) {
+    if (rhs.getValue() == 1)
+      return lhs();
+  } else if (auto rhs = operands[1].dyn_cast_or_null<SplatElementsAttr>()) {
+    if (rhs.getSplatValue<IntegerAttr>().getValue() == 1)
+      return lhs();
+  } // Assumes all tensors of all ones are splats.
+
   return overflowOrDiv0 ? Attribute() : result;
 }
 
 //===----------------------------------------------------------------------===//
 // SignedRemIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult SignedRemIOp::fold(ArrayRef<Attribute> operands) {
OpFoldResult UnsignedDivIOp::fold(ArrayRef<Attribute> operands) {
   bool div0 = false;
   auto result = constFoldBinaryOp<IntegerAttr>(operands, [&](APInt a, APInt b) {
     if (div0 || !b) {
       div0 = true;
       return a;
     }
     return a.udiv(b);
   });
+
+  // Fold out division by one.
+  if (auto rhs = operands[1].dyn_cast_or_null<IntegerAttr>()) {
+    if (rhs.getValue() == 1)
+      return lhs();
+  } else if (auto rhs = operands[1].dyn_cast_or_null<SplatElementsAttr>()) {
+    if (rhs.getSplatValue<IntegerAttr>().getValue() == 1)
+      return lhs();
+  } // Assumes all tensors of all ones are splats.
+
   return div0 ? Attribute() : result;
 }
 
 //===----------------------------------------------------------------------===//
 // UnsignedRemIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult UnsignedRemIOp::fold(ArrayRef<Attribute> operands) {

mlir/test/Transforms/canonicalize.mlir

func @remove_dead_else(%M : memref<100 x i32>) {
   return
 }
 // CHECK:      affine.if
 // CHECK-NEXT:   affine.for
 // CHECK-NEXT:     affine.load
 // CHECK-NEXT:   }
 // CHECK-NEXT: }
 // CHECK-NEXT: affine.load
+
+// -----
+
+// CHECK-LABEL: func @divi_signed_by_one
+// CHECK-SAME: %[[ARG:[a-zA-Z0-9]+]]
+func @divi_signed_by_one(%arg0: i32) -> (i32) {
+  %c1 = constant 1 : i32
+  %res = divi_signed %arg0, %c1 : i32
+  // CHECK: return %[[ARG]]
+  return %res : i32
+}
+
+// CHECK-LABEL: func @divi_unsigned_by_one
+// CHECK-SAME: %[[ARG:[a-zA-Z0-9]+]]
+func @divi_unsigned_by_one(%arg0: i32) -> (i32) {
+  %c1 = constant 1 : i32
+  %res = divi_unsigned %arg0, %c1 : i32
+  // CHECK: return %[[ARG]]
+  return %res : i32
+}
+
+// CHECK-LABEL: func @tensor_divi_signed_by_one
+// CHECK-SAME: %[[ARG:[a-zA-Z0-9]+]]
+func @tensor_divi_signed_by_one(%arg0: tensor<4x5xi32>) -> tensor<4x5xi32> {
+  %c1 = constant dense<1> : tensor<4x5xi32>
+  %res = divi_signed %arg0, %c1 : tensor<4x5xi32>
+  // CHECK: return %[[ARG]]
+  return %res : tensor<4x5xi32>
+}
+
+// CHECK-LABEL: func @tensor_divi_unsigned_by_one
+// CHECK-SAME: %[[ARG:[a-zA-Z0-9]+]]
+func @tensor_divi_unsigned_by_one(%arg0: tensor<4x5xi32>) -> tensor<4x5xi32> {
+  %c1 = constant dense<1> : tensor<4x5xi32>
+  %res = divi_unsigned %arg0, %c1 : tensor<4x5xi32>
+  // CHECK: return %[[ARG]]
+  return %res : tensor<4x5xi32>
+}

mlir/test/Transforms/parallel-loop-collapsing.mlir

 // CHECK:         [[C12:%.*]] = constant 12 : index
 // CHECK:         [[C13:%.*]] = constant 13 : index
 // CHECK:         [[C3:%.*]] = constant 3 : index
 // CHECK:         [[C0:%.*]] = constant 0 : index
 // CHECK:         [[C1:%.*]] = constant 1 : index
 // CHECK:         [[C2:%.*]] = constant 2 : index
 // CHECK:         loop.parallel ([[NEW_I0:%.*]], [[NEW_I1:%.*]], [[NEW_I2:%.*]]) = ([[C0]], [[C0]], [[C0]]) to ([[C2]], [[C1]], [[C1]]) step ([[C1]], [[C1]], [[C1]]) {
 // CHECK:           [[I0:%.*]] = remi_signed [[NEW_I0]], [[C2]] : index
-// CHECK:           [[I3_COUNT:%.*]] = divi_signed [[NEW_I0]], [[C1]] : index
-// CHECK:           [[I4_COUNT:%.*]] = divi_signed [[NEW_I1]], [[C1]] : index
-// CHECK:           [[VAL_16:%.*]] = muli [[I4_COUNT]], [[C13]] : index
+// CHECK:           [[VAL_16:%.*]] = muli [[NEW_I1]], [[C13]] : index
 // CHECK:           [[I4:%.*]] = addi [[VAL_16]], [[C12]] : index
-// CHECK:           [[VAL_18:%.*]] = muli [[I3_COUNT]], [[C10]] : index
+// CHECK:           [[VAL_18:%.*]] = muli [[NEW_I0]], [[C10]] : index
 // CHECK:           [[I3:%.*]] = addi [[VAL_18]], [[C9]] : index
 // CHECK:           [[VAL_20:%.*]] = muli [[NEW_I2]], [[C7]] : index
 // CHECK:           [[I2:%.*]] = addi [[VAL_20]], [[C6]] : index
 // CHECK:           "magic.op"([[I0]], [[C3]], [[I2]], [[I3]], [[I4]]) : (index, index, index, index, index) -> index
 // CHECK:           loop.yield
 // CHECK-NEXT:    }
 // CHECK-NEXT:    return