numpy empty

Hardcode84 · Feb 21, 2021 · Feb 20, 2021 · Feb 20, 2021 · Feb 20, 2021 · Feb 20, 2021
commit 48b275870ae6ae7b17d675f932e1cea4ef1e0494
diff --git a/mlir-compiler/mlir-compiler/src/py_linalg_resolver.cpp b/mlir-compiler/mlir-compiler/src/py_linalg_resolver.cpp
@@ -27,12 +27,17 @@ struct PyBuilderContext
 
 namespace
 {
-bool is_compatible_types(mlir::TypeRange types)
+bool is_compatible_type(mlir::Type type)
 {
-    return !types.empty() && llvm::all_of(types, [](mlir::Type t)
+    if (auto tuple_type = type.dyn_cast<mlir::TupleType>())
     {
-        return t.isIntOrFloat() || t.isa<mlir::RankedTensorType>();
-    });
+        return llvm::all_of(tuple_type, &is_compatible_type);
+    }
+    return type.isIntOrFloat() || type.isa<mlir::RankedTensorType>();
+}
+bool is_compatible_types(mlir::TypeRange types)
+{
+    return !types.empty() && llvm::all_of(types, &is_compatible_type);
 }
 
 size_t py_func_arg_count(py::handle signature, py::handle func)
@@ -59,14 +64,30 @@ auto unwrap_ssa_val(py::handle obj)
 
 auto unwrap_shape(mlir::Location loc, mlir::OpBuilder& builder, mlir::Value shape)
 {
-    auto type = shape.getType().cast<mlir::TupleType>();
-    llvm::SmallVector<mlir::Value, 8> ret(type.size());
-    for (auto it : llvm::enumerate(ret))
+    llvm::SmallVector<mlir::Value, 8> ret;
+    auto type = shape.getType();
+    auto index_cast = [&](mlir::Value val)->mlir::Value
+    {
+        if (!val.getType().isa<mlir::IndexType>())
+        {
+            return builder.create<plier::CastOp>(loc, builder.getIndexType(), val);
+        }
+        return val;
+    };
+    if (auto tuple_type = type.dyn_cast<mlir::TupleType>())
+    {
+        ret.resize(tuple_type.size());
+        for (auto it : llvm::enumerate(ret))
+        {
+            auto elem_type = tuple_type.getType(it.index());
+            auto index = builder.create<mlir::ConstantIndexOp>(loc, it.index());
+            auto dim = builder.create<plier::GetItemOp>(loc, elem_type, shape, index);
+            it.value() = index_cast(dim);
+        }
+    }
+    else
     {
-        auto elem_type = type.getType(it.index());
-        auto index = builder.create<mlir::ConstantIndexOp>(loc, it.index());
-        auto dim = builder.create<plier::GetItemOp>(loc, elem_type, shape, index);
-        it.value() = dim;
+        ret.emplace_back(index_cast(shape));
     }
     return ret;
 }
@@ -479,7 +500,7 @@ py::object len_impl(py::capsule /*context*/, py::capsule ssa_val)
     {
         return py::int_(tuple_type.size());
     }
-    return py::int_(0);
+    return py::int_(1);
 }
 
 void setup_py_var(pybind11::handle var)

diff --git a/numba/mlir/numpy/funcs.py b/numba/mlir/numpy/funcs.py
@@ -57,9 +57,14 @@ def body(a, b):
     return eltwise(builder, arg, body, builder.float64)
 
 @register_func('numpy.square', numpy.square)
-def quare_impl(builder, arg):
+def square_impl(builder, arg):
 
     def body(a, b):
         return a * a
 
     return eltwise(builder, arg, body)
+
+@register_func('numpy.empty', numpy.empty)
+def empty_impl(builder, shape):
+    # TODO: dtype
+    return builder.init_tensor(shape, builder.float64)
diff --git a/numba/mlir/tests/test_numpy.py b/numba/mlir/tests/test_numpy.py
@@ -209,5 +209,26 @@ def py_func(a, b):
         arr = np.array([0.0])
         assert_equal(py_func(arr, 5), jit_func(arr, 5))
 
+    def test_empty1(self):
+        def py_func(d):
+            a = np.empty(d)
+            for i in range(d):
+                a[i] = i
+            return a
+
+        jit_func = njit(py_func)
+        assert_equal(py_func(5), jit_func(5))
+
+    def test_empty2(self):
+        def py_func(d1, d2):
+            a = np.empty((d1, d2))
+            for i in range(d1):
+                for j in range(d2):
+                    a[i, j] = i + j * 10
+            return a
+
+        jit_func = njit(py_func)
+        assert_equal(py_func(5, 7), jit_func(5, 7))
+
 if __name__ == '__main__':
     unittest.main()