Skip to content

[MLIR][Linalg][Python] Improve bindings for linalg.elementwise #139462

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
May 12, 2025
Merged

[MLIR][Linalg][Python] Improve bindings for linalg.elementwise #139462

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
87f389d
[MLIR][Linalg][Python] Improve bindings for linalg.elementwise
rolfmorel May 11, 2025
4837a05
Minor improvements
rolfmorel May 12, 2025
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
61 changes: 61 additions & 0 deletions mlir/python/mlir/dialects/linalg/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -216,6 +216,67 @@ def contract(
)


# Extend and shadow the TableGen-derived version to make sure correct default
# indexing_maps are derived (as there is no mechanism for doing so given the
# Python API bypasses the C++-builders).
class ElementwiseOp_(ElementwiseOp):
def __init__(
self,
result_tensors,
inputs,
outputs,
kind,
*,
indexing_maps=None,
loc=None,
ip=None,
):
if indexing_maps is None:
inputs = [_get_op_result_or_value(in_) for in_ in inputs]
for in0, in1 in zip(inputs[:-1], inputs[1:]):
assert in0.type == in1.type
output = _get_op_result_or_value(outputs[0])
assert inputs[0].type == output.type
num_args = len(inputs) + 1
indexing_maps = [AffineMap.get_identity(output.type.rank)] * num_args

super().__init__(
result_tensors=result_tensors,
inputs=inputs,
outputs=outputs,
kind=kind,
indexing_maps=indexing_maps,
loc=loc,
ip=ip,
)


ElementwiseOp = ElementwiseOp_


def elementwise(
*ins: Union[Operation, OpView, Value],
outs: Sequence[Union[Operation, OpView, Value]],
kind: Union[ElementwiseKind, Attribute],
indexing_maps: Optional[Sequence[AffineMapAttr]] = None,
):
ins = [_get_op_result_or_value(input) for input in ins]
if len(outs) != 1:
raise ValueError(f"{outs=} must have length 1.")
init = _get_op_result_or_value(outs[0])
result_types = [init.type] if isinstance(init.type, RankedTensorType) else []

op = ElementwiseOp(
result_tensors=result_types,
inputs=ins,
outputs=[init],
kind=kind,
indexing_maps=indexing_maps,
)
fill_builtin_region(op.operation)
return _get_op_result_or_op_results(op)


def pack(
source,
dest,
Expand Down
186 changes: 186 additions & 0 deletions mlir/test/python/dialects/linalg/ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -606,3 +606,189 @@ def tensor_pack(src, dst):
# CHECK: return %[[VAL_4]] : tensor<128x128xf32>
# CHECK: }
print(module)


# CHECK-LABEL: TEST: testElementwiseOp
@run
def testElementwiseOp():
with Context(), Location.unknown():
module = Module.create()
f32 = F32Type.get()
with InsertionPoint(module.body):
rect_shape = (8, 16)
vert_line_shape = (8,)
hor_line_shape = (16,)
transposed_rect_shape = (16, 8)

# CHECK-DAG: #[[$IdentMap2D:.*]] = affine_map<(d0, d1) -> (d0, d1)>
# CHECK-DAG: #[[$TransMap2D:.*]] = affine_map<(d0, d1) -> (d1, d0)>
# CHECK-DAG: #[[$VertLineBCastMap:.*]] = affine_map<(d0, d1) -> (d0)>
# CHECK-DAG: #[[$HorLineBCastMap:.*]] = affine_map<(d0, d1) -> (d1)>

ident_map_2d = AffineMap.get_identity(2)
transposed_map_2d = AffineMap.get_permutation((1, 0))
vert_line_bcast_map = AffineMap.get(2, 0, [AffineDimExpr.get(0)])
hor_line_bcast_map = AffineMap.get(2, 0, [AffineDimExpr.get(1)])

# CHECK: func.func @elementwise_op(
@func.FuncOp.from_py_func(
# CHECK-SAME: %[[Rect:.*]]: tensor<8x16xf32>,
RankedTensorType.get(rect_shape, f32),
# CHECK-SAME: %[[RectMem:.*]]: memref<8x16xf32>,
MemRefType.get(rect_shape, f32),
# CHECK-SAME: %[[VertLine:.*]]: tensor<8xf32>,
RankedTensorType.get(vert_line_shape, f32),
# CHECK-SAME: %[[VertLineMem:.*]]: memref<8xf32>,
MemRefType.get(vert_line_shape, f32),
# CHECK-SAME: %[[HorLine:.*]]: tensor<16xf32>,
RankedTensorType.get(hor_line_shape, f32),
# CHECK-SAME: %[[HorLineMem:.*]]: memref<16xf32>,
MemRefType.get(hor_line_shape, f32),
# CHECK-SAME: %[[TransRect:.*]]: tensor<16x8xf32>,
RankedTensorType.get(transposed_rect_shape, f32),
# CHECK-SAME: %[[TransRectMem:.*]]: memref<16x8xf32>)
MemRefType.get(transposed_rect_shape, f32),
)
def elementwise_op(
rect,
rect_mem,
vert_line,
vert_line_mem,
hor_line,
hor_line_mem,
trans_rect,
trans_rect_mem,
):
# CHECK: %[[OutRect:.*]] = tensor.empty() : tensor<8x16xf32>
out_rect = tensor.EmptyOp(rect_shape, f32)
# CHECK: %[[OutRectMem:.*]] = memref.alloca() : memref<8x16xf32>
out_rect_mem = memref.alloca(MemRefType.get(rect_shape, f32), [], [])

if _inferred_affine_maps := True:
# CHECK: linalg.elementwise
# CHECK-SAME: kind=#linalg.elementwise_kind<exp>
# CHECK-SAME: ins(%[[Rect]] : tensor<8x16xf32>)
# CHECK-SAME: outs(%[[OutRect]] : tensor<8x16xf32>) -> tensor<8x16xf32>
op1 = linalg.ElementwiseOp(
result_tensors=(out_rect.result.type,),
inputs=(rect,),
outputs=(out_rect,),
kind=linalg.ElementwiseKind.exp,
)
linalg.fill_builtin_region(op1.operation)

# CHECK: linalg.elementwise
# CHECK-SAME: kind=#linalg.elementwise_kind<exp>
# CHECK-SAME: ins(%[[Rect]] : tensor<8x16xf32>)
# CHECK-SAME: outs(%[[OutRect]] : tensor<8x16xf32>) -> tensor<8x16xf32>
linalg.elementwise(
rect,
outs=(out_rect,),
kind=linalg.ElementwiseKind.exp,
)

# CHECK: linalg.elementwise
# CHECK-SAME: kind=#linalg.elementwise_kind<exp>
# CHECK-SAME: ins(%[[RectMem]] : memref<8x16xf32>)
# CHECK-SAME: outs(%[[OutRectMem]] : memref<8x16xf32>)
linalg.elementwise(
rect_mem,
outs=(out_rect_mem,),
kind=linalg.ElementwiseKind.exp,
)

if _explicit_ident_affine_maps := True:
# Same as above but with default identity indexing_maps explicitly provided.
# CHECK: linalg.elementwise
# CHECK-SAME: kind=#linalg.elementwise_kind<exp>
# CHECK-SAME: ins(%[[Rect]] : tensor<8x16xf32>)
# CHECK-SAME: outs(%[[OutRect]] : tensor<8x16xf32>) -> tensor<8x16xf32>
op3 = linalg.ElementwiseOp(
result_tensors=(out_rect.result.type,),
inputs=(rect,),
outputs=(out_rect,),
kind=linalg.ElementwiseKind.exp,
indexing_maps=[ident_map_2d, ident_map_2d],
)
linalg.fill_builtin_region(op3.operation)

# CHECK: linalg.elementwise
# CHECK-SAME: kind=#linalg.elementwise_kind<exp>
# CHECK-SAME: ins(%[[RectMem]] : memref<8x16xf32>)
# CHECK-SAME: outs(%[[OutRectMem]] : memref<8x16xf32>)
linalg.elementwise(
rect_mem,
outs=(out_rect_mem,),
kind=linalg.ElementwiseKind.exp,
indexing_maps=[ident_map_2d, ident_map_2d],
)

if _ops_with_non_ident_input_maps := True:
# CHECK: linalg.elementwise kind=#linalg.elementwise_kind<exp>
# CHECK-SAME: indexing_maps = [#[[$VertLineBCastMap]], #[[$IdentMap2D]]]
# CHECK-SAME: ins(%[[VertLine]] : tensor<8xf32>)
# CHECK-SAME: outs(%[[OutRect]] : tensor<8x16xf32>) -> tensor<8x16xf32>
op4 = linalg.ElementwiseOp(
result_tensors=(out_rect.result.type,),
inputs=(vert_line,),
outputs=(out_rect,),
kind=linalg.ElementwiseKind.exp,
indexing_maps=[vert_line_bcast_map, ident_map_2d],
)
linalg.fill_builtin_region(op4.operation)

# CHECK: linalg.elementwise kind=#linalg.elementwise_kind<add>
# CHECK-SAME: indexing_maps = [#[[$IdentMap2D]], #[[$VertLineBCastMap]], #[[$IdentMap2D]]]
# CHECK-SAME: ins(%[[Rect]], %[[VertLine]] : tensor<8x16xf32>, tensor<8xf32>)
# CHECK-SAME: outs(%[[OutRect]] : tensor<8x16xf32>) -> tensor<8x16xf32>
op4 = linalg.ElementwiseOp(
result_tensors=(out_rect.result.type,),
inputs=(rect, vert_line),
outputs=(out_rect,),
kind=linalg.ElementwiseKind.add,
indexing_maps=[ident_map_2d, vert_line_bcast_map, ident_map_2d],
)
linalg.fill_builtin_region(op4.operation)

# CHECK: linalg.elementwise kind=#linalg.elementwise_kind<div>
# CHECK-SAME: indexing_maps = [#[[$VertLineBCastMap]], #[[$HorLineBCastMap]], #[[$IdentMap2D]]]
# CHECK-SAME: ins(%[[VertLine]], %[[HorLine]] : tensor<8xf32>, tensor<16xf32>)
# CHECK-SAME: outs(%[[OutRect]] : tensor<8x16xf32>) -> tensor<8x16xf32>
linalg.elementwise(
vert_line,
hor_line,
outs=(out_rect,),
kind=linalg.ElementwiseKind.div,
indexing_maps=[
vert_line_bcast_map,
hor_line_bcast_map,
ident_map_2d,
],
)

if _ops_with_non_ident_and_transposed_input_maps := True:
# CHECK: %[[VertLineBoolsMem:.*]] = memref.alloca() : memref<8xi1>
vert_line_bools_mem = memref.alloca(
MemRefType.get(vert_line_shape, IntegerType.get_signless(1)),
[],
[],
)
# CHECK: linalg.elementwise kind=#linalg.elementwise_kind<select>
# CHECK-SAME: indexing_maps = [#[[$VertLineBCastMap]], #[[$HorLineBCastMap]], #[[$TransMap2D]], #[[$IdentMap2D]]]
# CHECK-SAME: ins(%[[VertLineBoolsMem]], %[[HorLineMem]], %[[TransRectMem]] : memref<8xi1>, memref<16xf32>, memref<16x8xf32>)
# CHECK-SAME: outs(%[[OutRectMem]] : memref<8x16xf32>)
linalg.elementwise(
vert_line_bools_mem,
hor_line_mem,
trans_rect_mem,
outs=(out_rect_mem,),
kind=linalg.ElementwiseKind.select,
indexing_maps=[
vert_line_bcast_map,
hor_line_bcast_map,
transposed_map_2d,
ident_map_2d,
],
)

print(module)