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
Next Next commit
fixes to broadcasting
  • Loading branch information
Butygin
Butygin committed Feb 27, 2021
commit afdb06845d045c74e010e216ec3a075c0e73eb01
14 changes: 7 additions & 7 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 @@ -473,8 +473,7 @@ py::object broadcast_impl(py::capsule context, py::tuple args)
}
}

llvm::SmallVector<int64_t, 8> shape(static_cast<size_t>(shape_vals.size()), -1);
py::tuple ret(args.size());
py::tuple ret(mlir_args.size());
if (shape_vals.empty())
{
for (auto it : llvm::enumerate(mlir_args))
Expand All @@ -489,6 +488,7 @@ py::object broadcast_impl(py::capsule context, py::tuple args)
return std::move(ret);
}

llvm::SmallVector<int64_t, 8> shape(static_cast<size_t>(shape_vals.size()), -1);
auto tensor_type = mlir::RankedTensorType::get(shape, res_type);
for (auto it : llvm::enumerate(mlir_args))
{
Expand All @@ -499,14 +499,14 @@ py::object broadcast_impl(py::capsule context, py::tuple args)
if (auto src_type = type.dyn_cast<mlir::ShapedType>())
{
assert(src_type.hasRank());
auto num_dims = static_cast<unsigned>(src_type.getRank());
auto src_num_dims = static_cast<unsigned>(src_type.getRank());
auto num_dims = static_cast<unsigned>(tensor_type.getRank());
auto init = builder.create<mlir::linalg::InitTensorOp>(loc, shape_vals, tensor_type.getElementType()).getResult();
llvm::SmallVector<llvm::StringRef, 8> iterators(num_dims, "parallel");
auto map = mlir::AffineMap::getMultiDimIdentityMap(num_dims, builder.getContext());
mlir::AffineMap maps[] = {
map,
map,
mlir::AffineMap::getMinorIdentityMap(num_dims, src_num_dims, builder.getContext()),
mlir::AffineMap::getMultiDimIdentityMap(num_dims, builder.getContext()),
};
llvm::SmallVector<llvm::StringRef, 8> iterators(num_dims, "parallel");
auto body = [&](mlir::OpBuilder &builder, mlir::Location loc, mlir::ValueRange values)
{
assert(values.size() == 2);
Expand Down
2 changes: 1 addition & 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])]
test_data = [1, 2.5, np.array([1,2,3]), np.array([4.4,5.5,6.6]), np.array([[7,8,9]])]
for py_func in funcs:
jit_func = njit(py_func)
for a1, a2 in itertools.product(test_data, test_data):
Expand Down