[Python Bindings] How to create a `tensor.extract_slice` operation with only static offsets, static sizes and static strides?

MLIR
1

Hi!
I want to use the tensor.extract_slice operation to slice a tensor. According to the documentation here, although this operation supports dynamic and static offsets, sizes and strides, someone could create a tensor.extract_slice operation with only static offsets, sizes and strides, just like the following example:

%1 = tensor.extract_slice %0[0, 0, 0][1, 16, 4][1, 1, 1] :
  tensor<8x16x4xf32> to tensor<16x4xf32>

However, the python binding of this operation specify that offsets, sizes, strides, static_offsets, static_sizes and static_strides are all positional argument. It means that some one need to fill all these arguments even if there are only static offsets, sizes and strides. To make things clearer, here is part of the __init__

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class ExtractSliceOp(_ods_ir. OpView):
  OPERATION_NAME = "tensor.extract_slice"

  _ODS_OPERAND_SEGMENTS = [1,-1,-1,-1,]

  _ODS_REGIONS = (0, True)

  def __init__(self, result, source, offsets, sizes, strides, static_offsets, static_sizes, static_strides, *, loc=None, ip=None):
    ...

Could anyone provide me some guidance on how to use the python binding to create a tensor.extract_slice operation with only static offsets, sizes and strides? Any insight would be greatly appreciated

2

I think just specifying empty lists or None for those should work. The generated builders are somewhat reasonable defaults that work well enough, but one can also add a more convenient builder (e.g., https://github.com/llvm/llvm-project/blob/main/mlir/python/mlir/dialects/_scf_ops_ext.py) on top for more ergonomic usage where they fail (in there you could for example consider based on input type if it is dynamic or static offset and so make a builder that would be rather tricky to autogenerate).

3

Thank you for your reply! If I understand correctly, are you suggesting that the parameters offsets, sizes, and strides should be set to either empty lists or None? I have give it a try:

  op = tensor.ExtractSliceOp(extract_slice_result_type, 
                             input_tensor,
                             None, None, None,
                             new_sizes_attr,
                             offsets_attr,
                             strides_attr)

However, I encountered the following error:

File "/root/project/llvm/build/tools/mlir/python_packages/mlir_core/mlir/dialects/_tensor_ops_gen.py", line 257, in __init__
  operands.append(_get_op_results_or_values(offsets))
File "/root/project/llvm/build/tools/mlir/python_packages/mlir_core/mlir/dialects/_ods_common.py", line 169, in get_op_results_or_values
  return [get_op_result_or_value(element) for element in arg]

If I’m not mistaken, the __init__ function will check whether the types of offsets, sizes, strides are Value or a list of Value through the get_op_results_or_values function. Could you please confirm if my understanding is correct, or if there’s anything important I might have missed?

4

here’s a convenience wrapper that should work

def extract_slice(
    source: "Tensor",
    offsets: Optional[Sequence[Value]] = None,
    strides: Optional[Sequence[Value]] = None,
    static_offsets: Optional[Sequence[int]] = None,
    static_sizes: Optional[Sequence[int]] = None,
    static_strides: Optional[Sequence[int]] = None,
    *,
    loc=None,
    ip=None,
):
    if offsets is None:
        offsets = []
    if strides is None:
        strides = []
    assert static_sizes, f"this convenience method only handles static sizes"
    assert offsets or static_offsets and bool(offsets) != bool(static_offsets)
    assert strides or static_strides and bool(strides) != bool(static_strides)
    sizes = []
    source_type = RankedTensorType(source.type)
    result = RankedTensorType.get(static_sizes, source_type.element_type)
    return tensor.ExtractSliceOp(
        result,
        source,
        offsets,
        sizes,
        strides,
        static_offsets,
        static_sizes,
        static_strides,
        loc=loc,
        ip=ip,
    )
5

If you want a fully static extract_slice, wouldn’t you write a wrapper that does not even have the offsets/strides parameters?

def static_extract_slice(
    source: "Tensor",
    offsets: Optional[Sequence[int]] = None,
    sizes: Optional[Sequence[int]] = None,
    strides: Optional[Sequence[int]] = None,
    *,
    loc=None,
    ip=None,
):
1 Like
6

sure but for me (i.e. from where I copy+pasted from) this wrapper is meant to handle (dynamic, static)x(offsets, strides). i just had it handy so i copy+pasted it here.

7

Thank you for the code snippet! It helps me a lot, I solve my problem by using empty list for the unwanted arguments :slight_smile: