work on linalg resolver

Hardcode84 · Feb 21, 2021 · Feb 20, 2021 · Feb 20, 2021 · Feb 20, 2021 · Feb 20, 2021
commit 8557f0fb74de07ea3df87bc6251f47e6d6976048
diff --git a/mlir-compiler/mlir-compiler/src/py_linalg_resolver.cpp b/mlir-compiler/mlir-compiler/src/py_linalg_resolver.cpp
@@ -354,7 +354,8 @@ py::object init_tensor_impl(py::capsule context, py::handle shape, py::capsule d
     auto& ctx = get_py_context(context);
     auto elem_type = unwrap_mlir<mlir::Type>(dtype);
     mlir::Value init;
-    if (py::len(shape) == 0)
+    auto count = py::len(shape);
+    if (count == 0)
     {
         // TODO: undef
         auto zero_val = plier::getZeroVal(elem_type);
@@ -363,9 +364,11 @@ py::object init_tensor_impl(py::capsule context, py::handle shape, py::capsule d
     }
     else
     {
-        auto& builder = ctx.builder;
-        auto loc = ctx.loc;
-        auto shape_val = unwrap_shape(loc, builder, unwrap_ssa_val(shape)) ;
+        llvm::SmallVector<mlir::Value, 8> shape_val(count);
+        for (size_t i = 0; i < count; ++i)
+        {
+            shape_val[i] = unwrap_ssa_val(shape[py::int_(i)]);
+        }
         init = ctx.builder.create<mlir::linalg::InitTensorOp>(ctx.loc, shape_val, elem_type);
     }
     return ctx.context.create_var(context, init);
@@ -571,11 +574,41 @@ py::object len_impl(py::capsule /*context*/, py::capsule ssa_val)
     return py::int_(1);
 }
 
+py::object getitem_impl(py::capsule context, py::capsule ssa_val, py::handle index)
+{
+    auto& ctx = get_py_context(context);
+    auto value = unwrap_mlir<mlir::Value>(ssa_val);
+    auto& builder = ctx.builder;
+    auto loc = ctx.loc;
+    auto index_val = index.cast<int64_t>();
+    auto type = value.getType();
+    if (auto tuple_type = type.dyn_cast<mlir::TupleType>())
+    {
+        if (index_val < 0 || index_val >= static_cast<int64_t>(tuple_type.size()))
+        {
+            plier::report_error("Invelid getitem index");
+        }
+        auto elem_type = tuple_type.getType(static_cast<size_t>(index_val));
+        auto ind = builder.create<mlir::ConstantIndexOp>(loc, index_val);
+        auto item = builder.create<plier::GetItemOp>(loc, elem_type, value, ind);
+        return ctx.context.create_var(context, item.getResult());
+    }
+    else
+    {
+        if (0 != index_val)
+        {
+            plier::report_error("Invelid getitem index");
+        }
+        return ctx.context.create_var(context, value);
+    }
+}
+
 void setup_py_var(pybind11::handle var)
 {
     py::setattr(var, "_shape", py::cpp_function(&shape_impl));
     py::setattr(var, "_dtype", py::cpp_function(&dtype_impl));
     py::setattr(var, "_len", py::cpp_function(&len_impl));
+    py::setattr(var, "_getitem", py::cpp_function(&getitem_impl));
 }
 
 PyLinalgResolver::Values unpack_results(py::handle object)
@@ -648,5 +681,9 @@ llvm::Optional<PyLinalgResolver::Values> PyLinalgResolver::rewrite(llvm::StringR
     setup_py_builder(py_builder, builder);
 
     auto result = builder_func(py_builder, *py_args);
+    if (result.is_none())
+    {
+        return {};
+    }
     return unpack_results(result);
 }
diff --git a/numba/mlir/linalg_builder.py b/numba/mlir/linalg_builder.py
@@ -16,6 +16,12 @@ def dtype(self):
     def __len__(self):
         return self._len(self._context, self._ssa_val)
 
+    def __getitem__(self, index):
+        return self._getitem(self._context, self._ssa_val, index)
+
+def is_literal(val):
+    return not isinstance(val, Var)
+
 class Builder:
     def __init__(self, context):
         self._context = context