[Mlir-commits] [mlir] [mlir][tosa] Fix intdiv folder type equality check (PR #180763)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Luke Hutton (lhutton1)

<details>
<summary>Changes</summary>

The folder was incorrectly checking the tensor types for equality as opposed to the element types.

cc @<!-- -->udaya-ranga

---
Full diff: https://github.com/llvm/llvm-project/pull/180763.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp (+4-1) 
- (modified) mlir/test/Dialect/Tosa/constant_folding.mlir (+5-5) 


``````````diff
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 42033ce8a3b02..a33ecb8511a92 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -1247,7 +1247,10 @@ OpFoldResult IntDivOp::fold(FoldAdaptor adaptor) {
   auto resultTy = llvm::dyn_cast<RankedTensorType>(getType());
   if (!lhsTy || !rhsTy || !resultTy)
     return {};
-  if (lhsTy != rhsTy)
+
+  const Type lhsETy = lhsTy.getElementType();
+  const Type rhsETy = rhsTy.getElementType();
+  if (lhsETy != rhsETy)
     return {};
 
   // IntDivOp inputs must be integer type, no need to check for quantized type
diff --git a/mlir/test/Dialect/Tosa/constant_folding.mlir b/mlir/test/Dialect/Tosa/constant_folding.mlir
index c3186279a30ae..56425e2feb0dd 100644
--- a/mlir/test/Dialect/Tosa/constant_folding.mlir
+++ b/mlir/test/Dialect/Tosa/constant_folding.mlir
@@ -158,13 +158,13 @@ func.func @fold_div_one_rhs_i32(%arg0: tensor<i32>) -> tensor<i32> {
 // -----
 
 // CHECK-LABEL: @fold_div_splat_i32
-func.func @fold_div_splat_i32() -> tensor<i32> {
-  %lhs = "tosa.const"() {values = dense<10> : tensor<i32>} : () -> tensor<i32>
-  %rhs = "tosa.const"() {values = dense<-3> : tensor<i32>} : () -> tensor<i32>
+func.func @fold_div_splat_i32() -> tensor<3x2xi32> {
+  %lhs = "tosa.const"() {values = dense<10> : tensor<1x2xi32>} : () -> tensor<1x2xi32>
+  %rhs = "tosa.const"() {values = dense<-3> : tensor<3x1xi32>} : () -> tensor<3x1xi32>
   // CHECK: %[[SPLAT:.+]] = "tosa.const"() <{values = dense<-3>
-  %div = tosa.intdiv %lhs, %rhs : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %div = tosa.intdiv %lhs, %rhs : (tensor<1x2xi32>, tensor<3x1xi32>) -> tensor<3x2xi32>
   // CHECK: return %[[SPLAT]]
-  return %div : tensor<i32>
+  return %div : tensor<3x2xi32>
 }
 
 // -----

``````````

</details>


https://github.com/llvm/llvm-project/pull/180763