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

[jreiffers]

Summary

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

[pifon2a]

Thank you. I wonder, why it was not added before...

mlir/include/mlir/Dialect/Arithmetic/Utils/Utils.h

 
 /// Helper struct to build simple arithmetic quantities with minimal type
 /// inference support.
 struct ArithBuilder {
   ArithBuilder(OpBuilder &b, Location loc) : b(b), loc(loc) {}
 
   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);
   Value slt(Value lhs, Value rhs);
 
 private:
   OpBuilder &b;
   Location loc;
 };
 } // namespace mlir
 
 #endif // MLIR_DIALECT_ARITHMETIC_UTILS_UTILS_H

mlir/lib/Dialect/Arithmetic/Utils/Utils.cpp

return llvm::to_vector<4>(
         return getValueOrCreateConstantIndexOp(b, loc, value);
       }));
 }
 
 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);
+  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);
+  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);
+  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);
+  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);
 }