IntelPython · Hardcode84 · Feb 27, 2021 · Feb 26, 2021 · Feb 26, 2021 · Feb 26, 2021
diff --git a/mlir-compiler/mlir-compiler/src/pipelines/plier_to_linalg.cpp b/mlir-compiler/mlir-compiler/src/pipelines/plier_to_linalg.cpp
@@ -214,12 +214,14 @@ mlir::LogicalResult lower_prange(plier::PyCallOp op, llvm::ArrayRef<mlir::Value>
 
 struct CallLowerer
 {
+    using args_t = llvm::ArrayRef<mlir::Value>;
+    using kwargs_t = llvm::ArrayRef<std::pair<llvm::StringRef, mlir::Value>>;
     mlir::LogicalResult operator()(
-        plier::PyCallOp op, llvm::StringRef name, llvm::ArrayRef<mlir::Value> args,
-        llvm::ArrayRef<std::pair<llvm::StringRef, mlir::Value>> kwargs,
+        plier::PyCallOp op, llvm::StringRef name, args_t args,
+        kwargs_t kwargs,
         mlir::PatternRewriter& rewriter)
     {
-        using func_t = mlir::LogicalResult(*)(plier::PyCallOp, llvm::ArrayRef<mlir::Value>, llvm::ArrayRef<std::pair<llvm::StringRef, mlir::Value>>, mlir::PatternRewriter&);
+        using func_t = mlir::LogicalResult(*)(plier::PyCallOp, args_t, kwargs_t, mlir::PatternRewriter&);
         std::pair<llvm::StringRef, func_t> handlers[] = {
             {"numba.prange", lower_prange},
         };
@@ -231,18 +233,8 @@ struct CallLowerer
             }
         }
 
-        if (auto result = linalg_resolver.rewrite_func(name, op.getLoc(), rewriter, args, kwargs))
+        if (mlir::succeeded(applyRewrite(op, rewriter, linalg_resolver.rewrite_func(name, op.getLoc(), rewriter, args, kwargs))))
         {
-            assert(result->size() == op->getNumResults());
-            rerun_std_pipeline(op);
-            if (result->empty())
-            {
-                rewriter.eraseOp(op);
-            }
-            else
-            {
-                rewriter.replaceOp(op, *result);
-            }
             return mlir::success();
         }
 
@@ -267,7 +259,39 @@ struct CallLowerer
             return mlir::failure();
         }
         auto full_name = (llvm::Twine("array.") + name).str();
-        if (auto result = linalg_resolver.rewrite_attr(full_name, op.getLoc(), rewriter, arg))
+        return applyRewrite(op, rewriter, linalg_resolver.rewrite_attr(full_name, op.getLoc(), rewriter, arg));
+    }
+
+    mlir::LogicalResult operator()(
+        plier::BinOp op, llvm::StringRef name, mlir::Value lhs, mlir::Value rhs,
+        mlir::PatternRewriter& rewriter)
+    {
+        if (!lhs.getType().isa<mlir::ShapedType>() &&
+            !rhs.getType().isa<mlir::ShapedType>())
+        {
+            return mlir::failure();
+        }
+        const std::pair<llvm::StringRef, llvm::StringRef> names[] = {
+            {"+", "operator.add"},
+            {"-", "operator.sub"},
+            {"*", "operator.mul"},
+        };
+        for (auto it : names)
+        {
+            if (it.first == name)
+            {
+                return applyRewrite(op, rewriter, linalg_resolver.rewrite_func(it.second, op.getLoc(), rewriter, {lhs, rhs}, {}));
+            }
+        }
+        return mlir::failure();
+    }
+
+private:
+    PyLinalgResolver linalg_resolver;
+
+    mlir::LogicalResult applyRewrite(mlir::Operation* op, mlir::PatternRewriter& rewriter, llvm::Optional<PyLinalgResolver::Values> result)
+    {
+        if (result)
         {
             assert(result->size() == op->getNumResults());
             rerun_std_pipeline(op);
@@ -283,9 +307,6 @@ struct CallLowerer
         }
         return mlir::failure();
     }
-
-private:
-    PyLinalgResolver linalg_resolver;
 };
 
 mlir::Value index_cast(mlir::Value value, mlir::Location loc, mlir::OpBuilder& builder)
@@ -659,8 +680,8 @@ struct RankedTypesCasts : public mlir::OpRewritePattern<plier::CastOp>
 
 struct GetattrRewriter : public mlir::OpRewritePattern<plier::GetattrOp>
 {
-    using resolver_t = llvm::function_ref<mlir::LogicalResult(plier::GetattrOp, llvm::StringRef, mlir::Value,
-                                                              mlir::PatternRewriter&)>;
+    using resolver_t = std::function<mlir::LogicalResult(plier::GetattrOp, llvm::StringRef, mlir::Value,
+                                                         mlir::PatternRewriter&)>;
 
     GetattrRewriter(mlir::TypeConverter &/*typeConverter*/,
                     mlir::MLIRContext *context,
@@ -679,6 +700,27 @@ struct GetattrRewriter : public mlir::OpRewritePattern<plier::GetattrOp>
     resolver_t resolver;
 };
 
+struct BinopRewriter : public mlir::OpRewritePattern<plier::BinOp>
+{
+    using resolver_t = std::function<mlir::LogicalResult(plier::BinOp, llvm::StringRef, mlir::Value, mlir::Value,
+                                                         mlir::PatternRewriter&)>;
+
+    BinopRewriter(mlir::TypeConverter &/*typeConverter*/,
+                  mlir::MLIRContext *context,
+                  resolver_t resolver):
+        OpRewritePattern(context),
+        resolver(resolver)
+    {}
+
+    mlir::LogicalResult matchAndRewrite(
+        plier::BinOp op, mlir::PatternRewriter &rewriter) const override
+    {
+        return resolver(op, op.op(), op.lhs(), op.rhs(), rewriter);
+    }
+
+private:
+    resolver_t resolver;
+};
 
 void PlierToLinalgPass::runOnOperation()
 {
@@ -710,8 +752,9 @@ void PlierToLinalgPass::runOnOperation()
 
     patterns.insert<
         plier::CallOpLowering,
-        GetattrRewriter
-        >(type_converter, context, callLowerer);
+        GetattrRewriter,
+        BinopRewriter
+        >(type_converter, context, std::ref(callLowerer));
 
     patterns.insert<
         GetitemOpLowering<plier::GetItemOp>,

diff --git a/mlir-compiler/mlir-compiler/src/pipelines/plier_to_std.cpp b/mlir-compiler/mlir-compiler/src/pipelines/plier_to_std.cpp
@@ -1404,7 +1404,7 @@ void PlierToStdPass::runOnOperation()
 
     patterns.insert<
         plier::CallOpLowering
-        >(type_converter, context, callLowerer);
+        >(type_converter, context, std::ref(callLowerer));
 
     mlir::populateStdExpandOpsPatterns(context, patterns);
 

diff --git a/mlir-compiler/plier/CMakeLists.txt b/mlir-compiler/plier/CMakeLists.txt
@@ -85,7 +85,7 @@ target_include_directories(${PLIER_LIB} PRIVATE
 
 target_include_directories(${PLIER_LIB} PUBLIC
     ./include
-    ${PROJECT_BINARY_DIR}/include
+    ${PROJECT_BINARY_DIR}/plier/include
     )
 
 add_dependencies(${PLIER_LIB} MLIRPlierOpsIncGen)
diff --git a/mlir-compiler/plier/include/plier/PlierOps.td b/mlir-compiler/plier/include/plier/PlierOps.td
@@ -56,14 +56,14 @@ def GlobalOp : Plier_Op<"global", [NoSideEffect]> {
 
 def BinOp : Plier_Op<"binop", []> {
     let arguments = (ins
-                     AnyType:$rhs,
                      AnyType:$lhs,
+                     AnyType:$rhs,
                      StrAttr:$op);
 
     let results = (outs AnyType);
 
     let builders = [
-        OpBuilderDAG<(ins "::mlir::Value":$rhs, "::mlir::Value":$lhs, "::mlir::StringRef ":$op)>
+        OpBuilderDAG<(ins "::mlir::Value":$lhs, "::mlir::Value":$rhs, "::mlir::StringRef ":$op)>
     ];
 }
 

diff --git a/mlir-compiler/plier/include/plier/rewrites/call_lowering.hpp b/mlir-compiler/plier/include/plier/rewrites/call_lowering.hpp
@@ -15,7 +15,7 @@ namespace plier
 {
 struct CallOpLowering : public mlir::OpRewritePattern<plier::PyCallOp>
 {
-    using resolver_t = llvm::function_ref<mlir::LogicalResult(plier::PyCallOp, llvm::StringRef, llvm::ArrayRef<mlir::Value>, llvm::ArrayRef<std::pair<llvm::StringRef, mlir::Value>> , mlir::PatternRewriter&)>;
+    using resolver_t = std::function<mlir::LogicalResult(plier::PyCallOp, llvm::StringRef, llvm::ArrayRef<mlir::Value>, llvm::ArrayRef<std::pair<llvm::StringRef, mlir::Value>> , mlir::PatternRewriter&)>;
 
     CallOpLowering(mlir::TypeConverter &typeConverter,
                    mlir::MLIRContext *context,

diff --git a/numba/mlir/numpy/funcs.py b/numba/mlir/numpy/funcs.py
@@ -30,12 +30,29 @@ def eltwise(builder, args, body, res_type = None):
     return builder.generic(args, init, iterators, maps, body)
 
 @register_func('numpy.add', numpy.add)
+@register_func('operator.add')
 def add_impl(builder, arg1, arg2):
     def body(a, b, c):
         return a + b
 
     return eltwise(builder, (arg1, arg2), body)
 
+@register_func('numpy.subtract', numpy.subtract)
+@register_func('operator.sub')
+def sub_impl(builder, arg1, arg2):
+    def body(a, b, c):
+        return a - b
+
+    return eltwise(builder, (arg1, arg2), body)
+
+@register_func('numpy.multiply', numpy.multiply)
+@register_func('operator.mul')
+def mul_impl(builder, arg1, arg2):
+    def body(a, b, c):
+        return a * b
+
+    return eltwise(builder, (arg1, arg2), body)
+
 @register_func('array.sum')
 @register_func('numpy.sum', numpy.sum)
 def sum_impl(builder, arg, axis=None):

diff --git a/numba/mlir/tests/test_numpy.py b/numba/mlir/tests/test_numpy.py
@@ -4,6 +4,7 @@
 import numpy as np
 from numba.tests.support import TestCase
 import unittest
+import itertools
 
 _arr_1d_int = [1,2,3,4,5,6,7,8]
 _arr_1d_float = [1.0,2.1,3.2,4.3,5.4,6.5,7.6,8.7]
@@ -46,8 +47,6 @@ def test_unary(self):
             lambda a: a.size,
             # lambda a: a.T, TODO: need fortran layout support
             lambda a: a.T.T,
-            lambda a: np.add(a, 1),
-            lambda a: np.add(a, 2.5),
         ]
 
         for py_func in funcs:
@@ -56,23 +55,21 @@ def test_unary(self):
                 arr = np.array(a)
                 assert_equal(py_func(arr), jit_func(arr))
 
-    def test_add(self):
-        def py_func(a, b):
-            return np.add(a, b)
-
-        jit_func = njit(py_func)
-        arr1 = np.array([1,2,3])
-        arr2 = np.array([4,5,6])
-        assert_equal(py_func(arr1,arr2), jit_func(arr1,arr2))
-
-    def test_add_scalar(self):
-        def py_func(a, b):
-            return np.add(a, b)
+    def test_binary(self):
+        funcs = [
+            lambda a, b: np.add(a, b),
+            lambda a, b: a + b,
+            lambda a, b: np.subtract(a, b),
+            lambda a, b: a - b,
+            lambda a, b: np.multiply(a, b),
+            lambda a, b: a * b,
+        ]
 
-        jit_func = njit(py_func)
-        arr1 = 1
-        arr2 = 2
-        assert_equal(py_func(arr1, arr2), jit_func(arr1, arr2))
+        test_data = [1, 2.5, np.array([1,2,3]), np.array([4.4,5.5,6.6])]
+        for py_func in funcs:
+            jit_func = njit(py_func)
+            for a1, a2 in itertools.product(test_data, test_data):
+                assert_equal(py_func(a1,a2), jit_func(a1,a2))
 
     def test_sum_axis(self):
         funcs = [