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

[MLIR] Numpy broadcasting #193

Merged
Hardcode84 merged 8 commits into from
Feb 28, 2021
Merged

[MLIR] Numpy broadcasting #193

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
afdb068
fixes to broadcasting
Feb 26, 2021
89dfda7
fix
Feb 26, 2021
78c8eac
work on broadcast
Feb 26, 2021
f5f30b7
some work on broadcast
Feb 26, 2021
41c2598
work on broadcast
Feb 26, 2021
4104a10
broadcasting
Feb 26, 2021
26e1b95
broadcast fix
Feb 26, 2021
0e5a12c
PostFusionOptPass
Feb 26, 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
work on broadcast
  • Loading branch information
Butygin
Butygin committed Feb 28, 2021
commit 41c2598552e7118bee5a1262725c6a04b9696de4
30 changes: 17 additions & 13 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 @@ -435,7 +435,11 @@ mlir::Value expand_dim(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value
auto cond = builder.create<mlir::CmpIOp>(loc, mlir::CmpIPredicate::eq, one, dim_val);
auto true_body = [&](mlir::OpBuilder &builder, mlir::Location loc)
{
auto init = builder.create<mlir::linalg::InitTensorOp>(loc, target_shape, src_type.getElementType()).getResult();
assert(dim < shape.size());
shape[dim] = 1;
mlir::Type casted_type = mlir::RankedTensorType::get(shape, src_type.getElementType());
auto casted = builder.create<mlir::tensor::CastOp>(loc, casted_type, src).getResult();
auto init = builder.create<mlir::linalg::InitTensorOp>(loc, target_shape.take_front(num_dims), src_type.getElementType()).getResult();
llvm::SmallVector<mlir::AffineExpr, 8> exprs(num_dims);
for (unsigned i = 0; i < num_dims; ++i)
{
Expand All @@ -460,7 +464,7 @@ mlir::Value expand_dim(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value
builder.create<mlir::linalg::YieldOp>(loc, values[0]);
};

auto res = builder.create<mlir::linalg::GenericOp>(loc, target_type, src, init, maps, iterators, body);
auto res = builder.create<mlir::linalg::GenericOp>(loc, target_type, casted, init, maps, iterators, body);
builder.create<mlir::scf::YieldOp>(loc, res.getResult(0));
};
auto false_body = [&](mlir::OpBuilder &builder, mlir::Location loc)
Expand All @@ -471,10 +475,9 @@ mlir::Value expand_dim(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value
return builder.create<mlir::scf::IfOp>(loc, target_type, cond, true_body, false_body).getResult(0);
}

mlir::Value expand_dims(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value val, mlir::ValueRange target_shape)
mlir::Value expand_dims(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value val, unsigned num_dims, mlir::ValueRange target_shape)
{
auto count = static_cast<unsigned>(target_shape.size());
for (unsigned i = 0; i < count; ++i)
for (unsigned i = 0; i < num_dims; ++i)
{
val = expand_dim(builder, loc, val, i, target_shape);
}
Expand Down Expand Up @@ -540,11 +543,11 @@ py::object broadcast_impl(py::capsule context, py::tuple args)
}
res_type = broadcast_type(res_type, shape_and_type->second);
auto new_shape_vals = shape_and_type->first;
for (auto it : llvm::zip(shape_vals, new_shape_vals))
for (auto it : llvm::enumerate(llvm::zip(shape_vals, new_shape_vals)))
{
auto& old_val = std::get<0>(it);
auto new_val = std::get<1>(it);
old_val = broadcast_dim(builder, loc, old_val, new_val);
auto old_val = std::get<0>(it.value());
auto new_val = std::get<1>(it.value());
shape_vals[it.index()] = broadcast_dim(builder, loc, old_val, new_val);
}

if (new_shape_vals.size() > shape_vals.size())
Expand Down Expand Up @@ -575,10 +578,11 @@ py::object broadcast_impl(py::capsule context, py::tuple args)
for (auto it : llvm::enumerate(mlir_args))
{
mlir::Value val = it.value();
// if (val.getType().isa<mlir::ShapedType>())
// {
// val = expand_dims(builder, loc, val, shape_vals);
// }
if (auto src_type = val.getType().dyn_cast<mlir::ShapedType>())
{
assert(src_type.hasRank());
val = expand_dims(builder, loc, val, static_cast<unsigned>(src_type.getRank()), shape_vals);
}
if (val.getType() != tensor_type)
{
auto type = val.getType();
Expand Down
14 changes: 13 additions & 1 deletion numba/mlir/tests/test_numpy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ def test_binary(self):
lambda a, b: a * b,
]

test_data = [1, 2.5, np.array([1,2,3]), np.array([4.4,5.5,6.6]), np.array([[7,8,9]])]
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):
Expand Down Expand Up @@ -276,5 +276,17 @@ def test_reshape(self):
for a in [arr1]:
assert_equal(py_func(a), jit_func(a))

def test_broadcast(self):
funcs = [
lambda a, b: np.add(a, b),
# lambda a, b: np.add(a, b.reshape((b.size, 1))),
]

a = np.array([[1,2],[3,4]])
b = np.array([5,6])
for py_func in funcs:
jit_func = njit(py_func)
assert_equal(py_func(a,b), jit_func(a,b))

if __name__ == '__main__':
unittest.main()