store context in Var

Hardcode84 · Feb 21, 2021 · Feb 20, 2021 · Feb 20, 2021 · Feb 20, 2021 · Feb 20, 2021
commit 5981ec01d7f765dd040932526e724618c1c55c0e
diff --git a/mlir-compiler/mlir-compiler/src/py_linalg_resolver.cpp b/mlir-compiler/mlir-compiler/src/py_linalg_resolver.cpp
@@ -125,7 +125,7 @@ struct PyLinalgResolver::Context
     py::handle compile_func;
     py::handle lookup_func;
 
-    py::object create_var(mlir::Location loc, mlir::OpBuilder& builder, mlir::Value value)
+    py::object create_var(mlir::Location loc, mlir::OpBuilder& builder, py::capsule context, mlir::Value value)
     {
         if (value.getType().isa<mlir::RankedTensorType>())
         {
@@ -142,9 +142,9 @@ struct PyLinalgResolver::Context
                 mlir::Value mlir_dim = builder.create<mlir::DimOp>(loc, value, it2.index());
                 py_shape[it2.index()] = make_dim_val(it2.value(), mlir_dim);
             }
-            return var(wrap_mlir(value), py_shape, wrap_mlir(elem_type));
+            return var(context, wrap_mlir(value), py_shape, wrap_mlir(elem_type));
         }
-        return var(wrap_mlir(value), py::list(), wrap_mlir(value.getType()));
+        return var(context, wrap_mlir(value), py::list(), wrap_mlir(value.getType()));
     }
 
     mlir::FuncOp compile_body(py::handle body, py::list arg_types)
@@ -156,20 +156,20 @@ struct PyLinalgResolver::Context
         return mlir_func;
     }
 
-    py::object wrap_result(mlir::Location loc, mlir::OpBuilder& builder, mlir::ValueRange values)
+    py::object wrap_result(mlir::Location loc, mlir::OpBuilder& builder, py::capsule context, mlir::ValueRange values)
     {
         if (values.empty())
         {
             return py::none();
         }
         if (values.size() == 1)
         {
-            return create_var(loc, builder, values.front());
+            return create_var(loc, builder, context, values.front());
         }
         py::tuple ret(values.size());
         for (auto it : llvm::enumerate(values))
         {
-            ret[it.index()] = create_var(loc, builder, it.value());
+            ret[it.index()] = create_var(loc, builder, context, it.value());
         }
         return std::move(ret);
     }
@@ -296,7 +296,7 @@ py::object init_tensor_impl(py::capsule context, py::list shape, py::capsule dty
     {
         init = ctx.builder.create<mlir::linalg::InitTensorOp>(ctx.loc, unwrap_shape(shape), elem_type);
     }
-    return ctx.context.create_var(ctx.loc, ctx.builder, init);
+    return ctx.context.create_var(ctx.loc, ctx.builder, context, init);
 }
 
 py::object generic_impl(py::capsule context, py::handle inputs, py::handle outputs, py::list iterators, py::list maps, py::handle body)
@@ -337,7 +337,7 @@ py::object generic_impl(py::capsule context, py::handle inputs, py::handle outpu
     {
         inputs_args.append(output_args.begin(), output_args.end());
         auto res = builder.create<mlir::CallOp>(loc, body_func, inputs_args);
-        return ctx.context.wrap_result(loc, builder, cast_values(res.getResults(), ret_types));
+        return ctx.context.wrap_result(loc, builder, context, cast_values(res.getResults(), ret_types));
     }
     else
     {
@@ -359,7 +359,7 @@ py::object generic_impl(py::capsule context, py::handle inputs, py::handle outpu
             affine_maps,
             mlir_iterators,
             body_builder);
-        return ctx.context.wrap_result(loc, builder, generic_op.getResults());
+        return ctx.context.wrap_result(loc, builder, context, generic_op.getResults());
     }
 }
 
@@ -400,7 +400,7 @@ py::object from_elements_impl(py::capsule context, py::handle values, py::capsul
         }
     });
     auto res = builder.create<mlir::tensor::FromElementsOp>(loc, vals);
-    return ctx.context.create_var(ctx.loc, ctx.builder, res);
+    return ctx.context.create_var(ctx.loc, ctx.builder, context, res);
 }
 
 py::object extract_impl(py::capsule context, py::handle value, py::handle indices)
@@ -426,7 +426,7 @@ py::object extract_impl(py::capsule context, py::handle value, py::handle indice
         }
     });
     auto res = builder.create<mlir::tensor::ExtractOp>(loc, get_var_value(value), ind);
-    return ctx.context.create_var(ctx.loc, ctx.builder, res);
+    return ctx.context.create_var(ctx.loc, ctx.builder, context, res);
 }
 
 void setup_py_builder(py::handle builder, mlir::OpBuilder& b)
@@ -507,7 +507,8 @@ llvm::Optional<PyLinalgResolver::Values> PyLinalgResolver::rewrite(llvm::StringR
     }
 
     PyBuilderContext py_builder_context{loc, builder, *context};
-    auto py_builder = context->builder(py::capsule(&py_builder_context));
+    auto py_context = py::capsule(&py_builder_context);
+    auto py_builder = context->builder(py_context);
     setup_py_builder(py_builder, builder);
 
     assert(!args.empty());
@@ -519,7 +520,7 @@ llvm::Optional<PyLinalgResolver::Values> PyLinalgResolver::rewrite(llvm::StringR
     {
         auto index = static_cast<unsigned>(it.index());
         auto mlir_arg = it.value();
-        py_args[index] = context->create_var(loc, builder, mlir_arg);
+        py_args[index] = context->create_var(loc, builder, py_context, mlir_arg);
     }
 
     auto result = builder_func(py_builder, *py_args);

diff --git a/numba/mlir/linalg_builder.py b/numba/mlir/linalg_builder.py
@@ -1,7 +1,8 @@
 from .func_registry import add_func
 
 class Var:
-    def __init__(self, ssa_val, shape, dtype):
+    def __init__(self, context, ssa_val, shape, dtype):
+        self._context = context
         self._ssa_val = ssa_val
         self._shape = shape
         self._dtype = dtype