How to call Python bindings?

39 views
Skip to first unread message

Axel Wong

unread,
May 28, 2024, 11:56:10 PMMay 28
to OpenXLA Discuss
Hi everyone,

I would like to use Python bindings stablehlo.reduce_window() with stablehlo.maximum() to create the following IR in text. 

%9 = "stablehlo.reduce_window"(%8, %4) ({
^bb0(%arg11: tensor<f32>, %arg12: tensor<f32>):
  %30 = stablehlo.maximum %arg11, %arg12 : tensor<f32>
  stablehlo.return %30 : tensor<f32>
}) {base_dilations = dense<1> : tensor<4xi64>, padding = dense<0> : tensor<4x2xi64>, window_dilations = dense<1> : tensor<4xi64>, window_dimensions = dense<[1, 1, 2, 2]> : tensor<4xi64>, window_strides = dense<[1, 1, 2, 2]> : tensor<4xi64>} : (tensor<1x6x28x28xf32>, tensor<f32>) -> tensor<1x6x14x14xf32>

How would the Python code looks like?
I would appreciate any help and ideas.

Thanks,
Axel

Peter Hawkins

unread,
May 29, 2024, 10:17:36 AMMay 29
to Axel Wong, OpenXLA Discuss

--
You received this message because you are subscribed to the Google Groups "OpenXLA Discuss" group.
To unsubscribe from this group and stop receiving emails from it, send an email to openxla-discu...@openxla.org.
To view this discussion on the web visit https://groups.google.com/a/openxla.org/d/msgid/openxla-discuss/611e1ae0-ff50-41fd-9e24-931c3ee1dd21n%40openxla.org.
For more options, visit https://groups.google.com/a/openxla.org/d/optout.

Axel Wong

unread,
May 29, 2024, 10:01:44 PMMay 29
to OpenXLA Discuss, Peter Hawkins, OpenXLA Discuss, Axel Wong
Thanks.
I've managed to find a solution as follows.

mp1 = stablehlo.reduce_window([RankedTensorType.get((1,6,14,14), F32Type.get())], [relu1], [vc4], ir.DenseI64ArrayAttr.get([1,1,2,2]),
                                    window_strides=ir.DenseI64ArrayAttr.get([1,1,2,2]),
                                    base_dilations=ir.DenseI64ArrayAttr.get([1,1,1,1]),
                                    window_dilations=ir.DenseI64ArrayAttr.get([1,1,1,1]),
                                    padding=ir.DenseI64ArrayAttr.get([0,0,0,0]))
block0 = mp1.owner.regions[0].blocks.append(RankedTensorType.get(tuple(), F32Type.get()), RankedTensorType.get(tuple(), F32Type.get()))
with InsertionPoint(block0):
          arg0, arg1 = block0.arguments
          m = stablehlo.maximum(arg0, arg1)
          stablehlo.return_([m])

Axel

Jacques Pienaar

unread,
May 29, 2024, 10:36:27 PMMay 29
to Axel Wong, OpenXLA Discuss, Peter Hawkins
The generated MLIR python ones allow one to define handlers for attribute types so that you should be able to do 

[1,1,2,2]

Instead of

ir.DenseI64ArrayAttr.get([1,1,2,2])

Well at least where there is no ambiguity (e.g., if both i32 and i64 were supported here).

See https://github.com/llvm/llvm-project/blob/30c10fda2ba539e70bff4f05625ec6358c0f7502/mlir/python/mlir/ir.py#L246 for how defined, nothing special or built in and one can do same for ones own attribute types or override default.

-- Jacques 


Reply all
Reply to author
Forward
0 new messages