[Mlir-commits] [mlir] a48f910 - [mlir][tosa] Added folders for tosa.div

[Rob Suderman]

Author: Rob Suderman
Date: 2022-08-29T17:02:43-07:00
New Revision: a48f910d7cd4dc12194c65afb2a51b7a50292f75

URL: https://github.com/llvm/llvm-project/commit/a48f910d7cd4dc12194c65afb2a51b7a50292f75
DIFF: https://github.com/llvm/llvm-project/commit/a48f910d7cd4dc12194c65afb2a51b7a50292f75.diff

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

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

Reviewed By: NatashaKnk

Differential Revision: https://reviews.llvm.org/D132693

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
    mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
    mlir/test/Dialect/Tosa/constant-op-fold.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index 29b40cdbcaaa..c19c9e4346d0 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -542,6 +542,8 @@ def Tosa_DivOp : Tosa_Op<"div", [
   let results = (outs
     Tosa_Int32Tensor:$output
   );
+
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 6e9765571e85..6af5b3d1017f 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -506,6 +506,40 @@ OpFoldResult AddOp::fold(ArrayRef<Attribute> operands) {
                                                             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 {};
+}
+
 namespace {
 DenseElementsAttr MulBinaryFolder(DenseElementsAttr lhs, DenseElementsAttr rhs,
                                   RankedTensorType ty, int32_t shift) {

diff  --git a/mlir/test/Dialect/Tosa/constant-op-fold.mlir b/mlir/test/Dialect/Tosa/constant-op-fold.mlir
index 5187143f9ee0..f5dad4698dc7 100644
--- a/mlir/test/Dialect/Tosa/constant-op-fold.mlir
+++ b/mlir/test/Dialect/Tosa/constant-op-fold.mlir
@@ -164,6 +164,39 @@ func.func @fold_add_splat_f32() -> 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: @fold_mul_zero_rhs_f32
 func.func @fold_mul_zero_rhs_f32(%arg0: tensor<f32>) -> tensor<f32> {