[Mlir-commits] [mlir] 34f4a9e - Add ArithBuilder::sub, make add, mul work with IndexTypes.

[Johannes Reifferscheid]

Author: Johannes Reifferscheid
Date: 2022-09-05T12:44:19+02:00
New Revision: 34f4a9ef2a1de4aad978d0185f77b484077849bf

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

LOG: Add ArithBuilder::sub, make add, mul work with IndexTypes.

sgt and slt already worked with IndexTypes, the others did not.

Reviewed By: pifon2a

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Arithmetic/Utils/Utils.h
    mlir/lib/Dialect/Arithmetic/Utils/Utils.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Arithmetic/Utils/Utils.h b/mlir/include/mlir/Dialect/Arithmetic/Utils/Utils.h
index 924de08e81af2..fd11bd3edcebf 100644
--- a/mlir/include/mlir/Dialect/Arithmetic/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/Arithmetic/Utils/Utils.h
@@ -99,6 +99,7 @@ struct ArithBuilder {
 
   Value _and(Value lhs, Value rhs);
   Value add(Value lhs, Value rhs);
+  Value sub(Value lhs, Value rhs);
   Value mul(Value lhs, Value rhs);
   Value select(Value cmp, Value lhs, Value rhs);
   Value sgt(Value lhs, Value rhs);

diff  --git a/mlir/lib/Dialect/Arithmetic/Utils/Utils.cpp b/mlir/lib/Dialect/Arithmetic/Utils/Utils.cpp
index b568891df66a1..ca61669fa7f2b 100644
--- a/mlir/lib/Dialect/Arithmetic/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Arithmetic/Utils/Utils.cpp
@@ -93,24 +93,29 @@ Value ArithBuilder::_and(Value lhs, Value rhs) {
   return b.create<arith::AndIOp>(loc, lhs, rhs);
 }
 Value ArithBuilder::add(Value lhs, Value rhs) {
-  if (lhs.getType().isa<IntegerType>())
-    return b.create<arith::AddIOp>(loc, lhs, rhs);
-  return b.create<arith::AddFOp>(loc, lhs, rhs);
+  if (lhs.getType().isa<FloatType>())
+    return b.create<arith::AddFOp>(loc, lhs, rhs);
+  return b.create<arith::AddIOp>(loc, lhs, rhs);
+}
+Value ArithBuilder::sub(Value lhs, Value rhs) {
+  if (lhs.getType().isa<FloatType>())
+    return b.create<arith::SubFOp>(loc, lhs, rhs);
+  return b.create<arith::SubIOp>(loc, lhs, rhs);
 }
 Value ArithBuilder::mul(Value lhs, Value rhs) {
-  if (lhs.getType().isa<IntegerType>())
-    return b.create<arith::MulIOp>(loc, lhs, rhs);
-  return b.create<arith::MulFOp>(loc, lhs, rhs);
+  if (lhs.getType().isa<FloatType>())
+    return b.create<arith::MulFOp>(loc, lhs, rhs);
+  return b.create<arith::MulIOp>(loc, lhs, rhs);
 }
 Value ArithBuilder::sgt(Value lhs, Value rhs) {
-  if (lhs.getType().isa<IndexType, IntegerType>())
-    return b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sgt, lhs, rhs);
-  return b.create<arith::CmpFOp>(loc, arith::CmpFPredicate::OGT, lhs, rhs);
+  if (lhs.getType().isa<FloatType>())
+    return b.create<arith::CmpFOp>(loc, arith::CmpFPredicate::OGT, lhs, rhs);
+  return b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sgt, lhs, rhs);
 }
 Value ArithBuilder::slt(Value lhs, Value rhs) {
-  if (lhs.getType().isa<IndexType, IntegerType>())
-    return b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, lhs, rhs);
-  return b.create<arith::CmpFOp>(loc, arith::CmpFPredicate::OLT, lhs, rhs);
+  if (lhs.getType().isa<FloatType>())
+    return b.create<arith::CmpFOp>(loc, arith::CmpFPredicate::OLT, lhs, rhs);
+  return b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, lhs, rhs);
 }
 Value ArithBuilder::select(Value cmp, Value lhs, Value rhs) {
   return b.create<arith::SelectOp>(loc, cmp, lhs, rhs);