[mlir][tosa] Added folders for tosa.div

[rsuderman]

Summary

Added folders for tosa.sub that handles bypassing divide by one,
and a zero numerator.

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

def Tosa_DivOp : Tosa_Op<"div", [
   let arguments = (ins
     Tosa_Int32Tensor:$input1,
     Tosa_Int32Tensor:$input2
   );
 
   let results = (outs
     Tosa_Int32Tensor:$output
   );
+
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: logical_and
 //===----------------------------------------------------------------------===//
 def Tosa_LogicalAndOp : Tosa_Op<"logical_and", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,

mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp

OpFoldResult AddOp::fold(ArrayRef<Attribute> operands) {
 
   if (!lhsAttr || !rhsAttr)
     return {};
 
   return BinaryFolder<std::plus<APInt>, std::plus<APFloat>>(lhsAttr, rhsAttr,
                                                             lhsTy);
 }
 
+OpFoldResult DivOp::fold(ArrayRef<Attribute> operands) {
+  auto lhsTy = getInput1().getType().dyn_cast<RankedTensorType>();
+  auto rhsTy = getInput2().getType().dyn_cast<RankedTensorType>();
+  auto resultTy = getType().dyn_cast<RankedTensorType>();
+  if (!lhsTy || !rhsTy || !resultTy)
+    return {};
+  if (lhsTy != rhsTy)
+    return {};
+
+  auto resultETy = resultTy.getElementType();
+  auto lhsAttr = operands[0].dyn_cast_or_null<DenseElementsAttr>();
+  auto rhsAttr = operands[1].dyn_cast_or_null<DenseElementsAttr>();
+  if (lhsAttr && lhsAttr.isSplat()) {
+    if (resultETy.isa<IntegerType>() && lhsAttr.getSplatValue<APInt>().isZero())
+      return lhsAttr;
+  }
+
+  if (rhsAttr && rhsAttr.isSplat()) {
+    if (resultETy.isa<IntegerType>() && rhsAttr.getSplatValue<APInt>().isOne())
+      return getInput1();
+  }
+
+  if (rhsAttr && lhsAttr && rhsAttr.isSplat() && lhsAttr.isSplat()) {
+    if (resultETy.isa<IntegerType>()) {
+      APInt l = lhsAttr.getSplatValue<APInt>();
+      APInt r = rhsAttr.getSplatValue<APInt>();
+      APInt result = l.sdiv(r);
+      return DenseElementsAttr::get(resultTy, result);
+    }
+  }
+
+  return {};
+}
+
 OpFoldResult CastOp::fold(ArrayRef<Attribute> operands) {
   if (getInput().getType() == getType())
     return getInput();
   return {};
 }
 
 OpFoldResult ConstOp::fold(ArrayRef<Attribute> operands) {
   assert(operands.empty() && "constant has no operands");

mlir/test/Dialect/Tosa/constant-op-fold.mlir

func.func @fold_add_splat_f32() -> tensor<10xf32> {
   %add = "tosa.add"(%one, %two) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
   // CHECK: %[[THREE:.+]] = "tosa.const"() {value = dense<3.000000e+00> : tensor<10xf32>}
   // CHECK: return %[[THREE]]
   return %add : tensor<10xf32>
 }
 
 // -----
 
+// CHECK-LABEL: @fold_div_zero_lhs_i32
+func.func @fold_div_zero_lhs_i32(%arg0: tensor<i32>) -> tensor<i32> {
+  %zero = "tosa.const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: %[[ZERO:.+]] = "tosa.const"() {value = dense<0> : tensor<i32>}
+  %div = "tosa.div"(%zero, %arg0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  // CHECK: return %[[ZERO]]
+  return %div : tensor<i32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_div_one_rhs_i32
+func.func @fold_div_one_rhs_i32(%arg0: tensor<i32>) -> tensor<i32> {
+  %one = "tosa.const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %div = "tosa.div"(%arg0, %one) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  // CHECK: return %arg0
+  return %div : tensor<i32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_div_splat_i32
+func.func @fold_div_splat_i32() -> tensor<i32> {
+  %lhs = "tosa.const"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
+  %rhs = "tosa.const"() {value = dense<-3> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: %[[SPLAT:.+]] = "tosa.const"() {value = dense<-3> : tensor<i32>}
+  %div = "tosa.div"(%lhs, %rhs) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  // CHECK: return %[[SPLAT]]
+  return %div : tensor<i32>
+}
+
+// -----
+
 // CHECK-LABEL: @slice_splat
 func.func @slice_splat() -> tensor<1x1x1xi32> {
   // CHECK: %[[SLICE:.+]] = "tosa.const"() {value = dense<42> : tensor<1x1x1xi32>}
   %splat = "tosa.const"() {value = dense<42> : tensor<4x5x6xi32>} : () -> tensor<4x5x6xi32>
   %slice = "tosa.slice"(%splat) { size = [1, 1, 1], start = [1, 2, 3] } : (tensor<4x5x6xi32>) -> tensor<1x1x1xi32>
   // CHECK: return %[[SLICE]]
   return %slice : tensor<1x1x1xi32>
 }