Add DivOp to the Shape dialect

[jpienaar]

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

144

Divides

149–152

I think lets turn these around and say it follows floor(x / y) [I mean the equation of interest to me is mod(a, n) = a - n * floor(a / n) and yes it corresponds to python operator but the motivation is inverted for me :) ]

mlir/lib/Dialect/Shape/IR/Shape.cpp

605

Same here

mlir/test/Dialect/Shape/canonicalize.mlir

889

Could you add div(10, -4) too? (wouldn't matter for div but could for mod, and then I'd want to just extend these tests in that case)

[Jing]

Resolved review comments.

[Jing]

Could you help me land the path?

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

def Shape_ConstSizeOp : Shape_Op<"const_size", [
   let results = (outs Shape_SizeType:$result);
 
   let builders = [OpBuilderDAG<(ins "int64_t":$value)>];
 
   let assemblyFormat = "$value attr-dict";
   let hasFolder = 1;
 }
 
+def Shape_DivOp : Shape_Op<"div", [NoSideEffect]> {
+  let summary = "Division of sizes and indices";
+  let description = [{
+    Divide two sizes or indices. If either operand is an error it will be
+    propagated to the result. The operands can be of type `size` or `index`. If
+    at least one of the operands can hold an error, i.e. if it is of type `size`,
+    the result must be of type `size`. If error propagation is not possible
+    because both operands are of type `index` then the result may be of type
+    `size` or `index`. If both operands and result are of type `index`, their
+    runtime values could be negative, and the division semantic follows the
+    integer division operator "//" in Python (e.g. rounding toward negative
+    infinity). If any of the values is of type `size`, the behavior for
+    negative value is undefined.
+  }];
+
+  let arguments = (ins Shape_SizeOrIndexType:$lhs,
+                       Shape_SizeOrIndexType:$rhs);
+  let results = (outs Shape_SizeOrIndexType:$result);
+
+  let assemblyFormat = [{
+    $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)
+  }];
+
+  let verifier = [{ return ::verifySizeOrIndexOp(*this); }];
+  let hasFolder = 1;
+}
+
 def Shape_ShapeEqOp : Shape_Op<"shape_eq", [Commutative, NoSideEffect]> {
   let summary = "Returns whether the input shapes or extent tensors are equal";
   let description = [{
     Takes two shape or extent tensor operands and determines whether they are
     equal. When extent tensors are compared to shapes they are regarded as their
     equivalent non-error shapes. Error shapes can be tested for equality like
     any other shape value, meaning that the error value is equal to itself.
   }];

mlir/lib/Dialect/Shape/IR/Shape.cpp

 // CstrRequireOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult CstrRequireOp::fold(ArrayRef<Attribute> operands) {
   return operands[0];
 }
 
 //===----------------------------------------------------------------------===//
+// DivOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult DivOp::fold(ArrayRef<Attribute> operands) {
+  auto lhs = operands[0].dyn_cast_or_null<IntegerAttr>();
+  if (!lhs)
+    return nullptr;
+  auto rhs = operands[1].dyn_cast_or_null<IntegerAttr>();
+  if (!rhs)
+    return nullptr;
+
+  // APInt does not follow python semantics when the result is negative.
+  // Rather, APInt rounds toward zero while python rounds toward negative
+  // infinity.
+  APInt quotient, remainder;
+  APInt::sdivrem(lhs.getValue(), rhs.getValue(), quotient, remainder);
+  if (quotient.isNegative() && !remainder.isNullValue()) {
+    quotient -= 1;
+  }
+
+  Type indexTy = IndexType::get(getContext());
+  return IntegerAttr::get(indexTy, quotient);
+}
+
+//===----------------------------------------------------------------------===//
 // ShapeEqOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ShapeEqOp::fold(ArrayRef<Attribute> operands) {
   if (lhs() == rhs())
     return BoolAttr::get(getContext(), true);
   auto lhs = operands[0].dyn_cast_or_null<DenseIntElementsAttr>();
   if (lhs == nullptr)

mlir/test/Dialect/Shape/canonicalize.mlir

func @fold_mul_mixed() -> !shape.size {
   %c2 = shape.const_size 2
   %c3 = constant 3 : index
   %result = shape.mul %c2, %c3 : !shape.size, index -> !shape.size
   return %result : !shape.size
 }
 
 // -----
 
+// Fold `div` for constant sizes.
+// CHECK-LABEL: @fold_div_size
+func @fold_div_size() -> !shape.size {
+  // CHECK: %[[RESULT:.*]] = shape.const_size 3
+  // CHECK: return %[[RESULT]] : !shape.size
+  %c2 = shape.const_size 10
+  %c3 = shape.const_size 3
+  %result = shape.div %c2, %c3 : !shape.size, !shape.size -> !shape.size
+  return %result : !shape.size
+}
+
+// -----
+
+// Fold `div` for constant indices.
+// CHECK-LABEL: @fold_div_index
+func @fold_div_index() -> index {
+  // CHECK: %[[RESULT:.*]] = constant 2 : index
+  // CHECK: return %[[RESULT]] : index
+  %c2 = constant 10 : index
+  %c3 = constant 4 : index
+  %result = shape.div %c2, %c3 : index, index -> index
+  return %result : index
+}
+
+// -----
+
+// Fold `div` for constant indices.
+// CHECK-LABEL: @fold_div_index_neg_lhs
+func @fold_div_index_neg_lhs() -> index {
+  // CHECK: %[[RESULT:.*]] = constant -3 : index
+  // CHECK: return %[[RESULT]] : index
+  %c2 = constant -10 : index
+  %c3 = constant 4 : index
+  %result = shape.div %c2, %c3 : index, index -> index
+  return %result : index
+}
+
+// -----
+
+// Fold `div` for mixed constants.
+// CHECK-LABEL: @fold_div_mixed
+func @fold_div_mixed() -> !shape.size {
+  // CHECK: %[[RESULT:.*]] = shape.const_size 4
+  // CHECK: return %[[RESULT]] : !shape.size
+  %c2 = shape.const_size 12
+  %c3 = constant 3 : index
+  %result = shape.div %c2, %c3 : !shape.size, index -> !shape.size
+  return %result : !shape.size
+}
+
+// -----
+
 // Fold index_cast when already on index.
 // CHECK-LABEL: @fold_index_cast_on_index
 func @fold_index_cast_on_index(%arg: index) -> index {
   // CHECK-NOT: size_to_index
   %casted = shape.size_to_index %arg : index
   return %casted : index
 }
 

mlir/test/Dialect/Shape/ops.mlir

func @mul(%size_arg : !shape.size, %index_arg : index) {
   %size_prod = shape.mul %size_arg, %size_arg
       : !shape.size, !shape.size -> !shape.size
   %index_prod = shape.mul %index_arg, %index_arg : index, index -> index
   %mixed_prod = shape.mul %size_arg, %index_arg
       : !shape.size, index -> !shape.size
   return
 }
 
+func @div(%size_arg : !shape.size, %index_arg : index) {
+  %size_div = shape.div %size_arg, %size_arg
+      : !shape.size, !shape.size -> !shape.size
+  %index_div = shape.div %index_arg, %index_arg : index, index -> index
+  %mixed_div = shape.div %size_arg, %index_arg
+      : !shape.size, index -> !shape.size
+  return
+}
+
 func @add(%size_arg : !shape.size, %index_arg : index) {
   %size_sum = shape.add %size_arg, %size_arg
       : !shape.size, !shape.size -> !shape.size
   %index_sum = shape.add %index_arg, %index_arg : index, index -> index
   %mixed_sum = shape.add %size_arg, %index_arg
       : !shape.size, index -> !shape.size
   return
 }