Add builtin function abs and numpy.abs (#178)

Author: Hardcode84

mlir/lib/transforms/uplift_math.cpp

@@ -90,6 +90,33 @@ struct UpliftMathCalls : public mlir::OpRewritePattern<mlir::func::CallOp> {
   }
 };
 
+struct UpliftFabsCalls : public mlir::OpRewritePattern<mlir::func::CallOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mlir::func::CallOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+    auto funcName = op.getCallee();
+    if (funcName.empty())
+      return mlir::failure();
+
+    if (funcName != "fabs" && funcName != "fabsf")
+      return mlir::failure();
+
+    auto isNotValidType = [](mlir::Type t) {
+      return !t.isa<mlir::FloatType>();
+    };
+
+    if (op.getNumResults() != 1 || op.getNumOperands() != 1 ||
+        llvm::any_of(op.getOperandTypes(), isNotValidType) ||
+        llvm::any_of(op.getResultTypes(), isNotValidType))
+      return mlir::failure();
+
+    rewriter.replaceOpWithNewOp<mlir::math::AbsOp>(op, op.operands()[0]);
+    return mlir::success();
+  }
+};
+
 struct UpliftFma : public mlir::OpRewritePattern<mlir::arith::AddFOp> {
   using OpRewritePattern::OpRewritePattern;
 
@@ -123,7 +150,7 @@ struct UpliftMathPass
           plier::DependentDialectsList<mlir::func::FuncDialect,
                                        mlir::arith::ArithmeticDialect,
                                        mlir::math::MathDialect>,
-          UpliftMathCalls, UpliftFma> {};
+          UpliftMathCalls, UpliftFabsCalls, UpliftFma> {};
 } // namespace
 
 void plier::populateUpliftmathPatterns(mlir::MLIRContext &context,

numba_dpcomp/numba_dpcomp/mlir/builtin/funcs.py

@@ -135,3 +135,18 @@ def func(builder, *args):
 
 _gen_math_funcs()
 del _gen_math_funcs
+
+
+@register_func("abs", abs)
+def abs_impl(builder, arg):
+    t = arg.type
+    if is_int(t, builder):
+        c = arg < 0
+        return builder.select(c, -arg, arg)
+    if is_float(t, builder):
+        fname = "fabs"
+        if t == builder.float32:
+            fname = fname + "f"
+
+        res = builder.cast(0, t)
+        return builder.external_call(fname, arg, res, decorate=False)

numba_dpcomp/numba_dpcomp/mlir/linalg_builder.py

@@ -83,6 +83,12 @@ def __eq__(self, o):
     def __ne__(self, o):
         return self._binop(self._context, self._ssa_val, o, "ne")
 
+    def __neg__(self):
+        return self._unop(self._context, self._ssa_val, "-")
+
+    def __pos__(self):
+        return self._unop(self._context, self._ssa_val, "+")
+
     def __str__(self):
         return self._str(self._context, self._ssa_val)
 

numba_dpcomp/numba_dpcomp/mlir/numpy/funcs.py

@@ -206,6 +206,7 @@ def bool_type(builder, t):
         (register_func("numpy.cos", numpy.cos), f64_type, lambda a, b: math.cos(a)),
         (register_func("numpy.exp", numpy.exp), f64_type, lambda a, b: math.exp(a)),
         (register_func("numpy.tanh", numpy.tanh), f64_type, lambda a, b: math.tanh(a)),
+        (register_func("numpy.abs", numpy.abs), None, lambda a, b: abs(a)),
         (
             register_func("numpy.logical_not", numpy.logical_not),
             bool_type,

numba_dpcomp/numba_dpcomp/mlir/tests/test_basic.py

@@ -100,7 +100,7 @@ def py_func(a, b):
 
 
 @parametrize_function_variants(
-    "py_func", ["lambda a: +a", "lambda a: -a", "lambda a: ~a",]
+    "py_func", ["lambda a: +a", "lambda a: -a", "lambda a: ~a", "lambda a: abs(a)",]
 )
 @pytest.mark.parametrize("val", _test_values)
 def test_unary_ops(py_func, val, request):
@@ -184,7 +184,7 @@ def test_cast(py_func, val):
 
 
 @pytest.mark.parametrize("val", [1, 5, 5.5])
-@pytest.mark.parametrize("name", ["sqrt", "log", "exp", "sin", "cos", "erf", "tanh",])
+@pytest.mark.parametrize("name", ["sqrt", "log", "exp", "sin", "cos", "erf", "tanh"])
 def test_math_uplifting1(val, name):
     py_func = eval(f"lambda a: math.{name}(a)")
 

numba_dpcomp/numba_dpcomp/mlir/tests/test_numpy.py

@@ -87,6 +87,8 @@ def _vectorize_reference(func, arg1):
         "lambda a: np.cos(a)",
         "lambda a: np.exp(a)",
         "lambda a: np.tanh(a)",
+        "lambda a: np.abs(a)",
+        "lambda a: np.absolute(a)",
         "lambda a: a.size",
         "lambda a: a.T",
         "lambda a: a.T.T",

numba_dpcomp/numba_dpcomp/mlir_compiler/lib/py_linalg_resolver.cpp

@@ -1742,6 +1742,38 @@ static py::object binopImpl(py::capsule context, py::capsule ssaVal,
   plier::reportError("Unhandled binop type");
 }
 
+static py::object unopImpl(py::capsule context, py::capsule ssaVal,
+                           py::str op) {
+  auto &ctx = getPyContext(context);
+  auto &builder = ctx.builder;
+  auto loc = ctx.loc;
+  auto val = unwrapMlir<mlir::Value>(ssaVal);
+
+  auto type = val.getType();
+  if (!type.isa<mlir::IntegerType, mlir::IndexType, mlir::FloatType>())
+    plier::reportError("Invalid unop arg type");
+
+  auto opName = static_cast<std::string>(op);
+  mlir::Value res;
+  if (opName == "+") {
+    res = val;
+  } else if (opName == "-") {
+    if (type.isa<mlir::FloatType>()) {
+      res = builder.create<mlir::arith::NegFOp>(loc, val);
+    } else {
+      auto signlessType = makeSignlessType(type);
+      auto zero = builder.getIntegerAttr(signlessType, 0);
+      auto zeroVal = builder.create<mlir::arith::ConstantOp>(loc, zero);
+      val = doSignCast(builder, loc, val);
+      res = builder.create<mlir::arith::SubIOp>(loc, zeroVal, val);
+      res = doSignCast(builder, loc, res, type);
+    }
+  } else {
+    plier::reportError("Unhandled unop type");
+  }
+  return ctx.context.createVar(context, res);
+}
+
 static py::object strImpl(py::capsule /*context*/, py::capsule ssaVal) {
   return py::str("Var: \"" + toStr(unwrapMlir<mlir::Value>(ssaVal)) + "\"");
 }
@@ -1753,6 +1785,7 @@ static void setupPyVar(pybind11::handle var) {
   py::setattr(var, "_len", py::cpp_function(&lenImpl));
   py::setattr(var, "_getitem", py::cpp_function(&getitemImpl));
   py::setattr(var, "_binop", py::cpp_function(&binopImpl));
+  py::setattr(var, "_unop", py::cpp_function(&unopImpl));
   py::setattr(var, "_str", py::cpp_function(&strImpl));
 }