Skip to content
This repository was archived by the owner on Jan 25, 2023. It is now read-only.

[MLIR] Numpy empty and sum axis #187

Merged
Hardcode84 merged 15 commits into from
Feb 21, 2021
Merged

[MLIR] Numpy empty and sum axis #187

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
5981ec0
store context in Var
Feb 20, 2021
d53b90f
rework shape accessor
Feb 20, 2021
71b35c4
dtype accessor rework
Feb 20, 2021
437ac11
remove unused code
Feb 20, 2021
24fb1b5
refactor shape
Feb 20, 2021
ba50774
remove unused code
Feb 20, 2021
42a0964
fix setitem lowering
Feb 20, 2021
48b2758
numpy empty
Feb 20, 2021
b3a2105
numpy.sum
Feb 20, 2021
1bec4ad
some kwargs support
Feb 20, 2021
5712278
linlag resolver kwargs support
Feb 20, 2021
75dfc8b
linalg resolver some literal support
Feb 20, 2021
8557f0f
work on linalg resolver
Feb 20, 2021
a1bb470
add symbolDCE pass
Feb 20, 2021
1c25bf6
numpy sum axis support
Feb 20, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
refactor shape
  • Loading branch information
Butygin
Butygin committed Feb 20, 2021
commit 24fb1b5d411419c92bfa9bc6d57aaa0f2a7f9afc
52 changes: 35 additions & 17 deletions mlir-compiler/mlir-compiler/src/py_linalg_resolver.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,13 +66,16 @@ auto unwrap_ssa_val(py::handle obj)
return unwrap_mlir<mlir::Value>(obj.attr("_ssa_val").cast<py::capsule>());
}

auto unwrap_shape(py::list shape)
auto unwrap_shape(mlir::Location loc, mlir::OpBuilder& builder, mlir::Value shape)
{
llvm::SmallVector<mlir::Value, 8> ret;
ret.reserve(shape.size());
for (auto elem : shape)
auto type = shape.getType().cast<mlir::TupleType>();
llvm::SmallVector<mlir::Value, 8> ret(type.size());
for (auto it : llvm::enumerate(ret))
{
ret.push_back(unwrap_ssa_val(elem));
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;
}
return ret;
}
Expand Down Expand Up @@ -267,12 +270,12 @@ py::object broadcast_impl(py::capsule /*context*/, py::tuple args)
}
}

py::object init_tensor_impl(py::capsule context, py::list shape, py::capsule dtype)
py::object init_tensor_impl(py::capsule context, py::handle shape, py::capsule dtype)
{
auto& ctx = get_py_context(context);
auto elem_type = unwrap_mlir<mlir::Type>(dtype);
mlir::Value init;
if (shape.empty())
if (py::len(shape) == 0)
{
// TODO: undef
auto zero_val = plier::getZeroVal(elem_type);
Expand All @@ -281,7 +284,10 @@ py::object init_tensor_impl(py::capsule context, py::list shape, py::capsule dty
}
else
{
init = ctx.builder.create<mlir::linalg::InitTensorOp>(ctx.loc, unwrap_shape(shape), elem_type);
auto& builder = ctx.builder;
auto loc = ctx.loc;
auto shape_val = unwrap_shape(loc, builder, unwrap_ssa_val(shape)) ;
init = ctx.builder.create<mlir::linalg::InitTensorOp>(ctx.loc, shape_val, elem_type);
}
return ctx.context.create_var(context, init);
}
Expand Down Expand Up @@ -447,19 +453,19 @@ py::object shape_impl(py::capsule context, py::capsule ssa_val)
{
auto& builder = ctx.builder;
auto loc = ctx.loc;
auto make_dim_val = [&](auto dim, auto ssa_val)
{
return ctx.context.val(get_dim(dim), wrap_mlir(ssa_val));
};
auto mlir_type = value.getType().cast<mlir::RankedTensorType>();
auto shape = mlir_type.getShape();
py::list py_shape(shape.size());
for (auto it2 : llvm::enumerate(shape))
llvm::SmallVector<mlir::Value, 8> shape_vals(shape.size());
for (auto it : llvm::enumerate(shape))
{
mlir::Value mlir_dim = builder.create<mlir::DimOp>(loc, value, it2.index());
py_shape[it2.index()] = make_dim_val(it2.value(), mlir_dim);
auto i = it.index();
mlir::Value mlir_dim = builder.create<mlir::DimOp>(loc, value, i);
shape_vals[i] = mlir_dim;
}
return std::move(py_shape);
llvm::SmallVector<mlir::Type, 8> shape_types(shape.size(), builder.getIndexType());
auto shape_type = mlir::TupleType::get(builder.getContext(), shape_types);
auto shape_var = builder.create<plier::BuildTupleOp>(loc, shape_type, shape_vals);
return ctx.context.create_var(context, shape_var.getResult());
}
return py::list();
}
Expand All @@ -475,10 +481,22 @@ py::object dtype_impl(py::capsule /*context*/, py::capsule ssa_val)
return wrap_mlir(type);
}

py::object len_impl(py::capsule /*context*/, py::capsule ssa_val)
{
auto value = unwrap_mlir<mlir::Value>(ssa_val);
auto type = value.getType();
if (auto tuple_type = type.dyn_cast<mlir::TupleType>())
{
return py::int_(tuple_type.size());
}
return py::int_(0);
}

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));
}

PyLinalgResolver::Values unpack_results(py::handle object)
Expand Down
3 changes: 3 additions & 0 deletions numba/mlir/linalg_builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,9 @@ def shape(self):
def dtype(self):
return self._dtype(self._context, self._ssa_val)

def __len__(self):
return self._len(self._context, self._ssa_val)



class Val:
Expand Down