if pooling on H and W, stride: len(stride) need to be equal to len(inp.shape) for NCHW, should be (1, 1, stride_h, stride_w) for NHWC, should be (1, stride_h, stride_w, 1) kernel: similar to stride, len(kernel) also need to be equal to len(inp.shape)
(
reducer,
unit_factory,
inp,
stride,
kernel,
padding,
base_dilation=None,
kernel_dilation=None,
oshape=None,
)
| 346 | |
| 347 | |
| 348 | def _pooling( |
| 349 | reducer, |
| 350 | unit_factory, |
| 351 | inp, |
| 352 | stride, |
| 353 | kernel, |
| 354 | padding, |
| 355 | base_dilation=None, |
| 356 | kernel_dilation=None, |
| 357 | oshape=None, |
| 358 | ): |
| 359 | """ |
| 360 | if pooling on H and W, |
| 361 | stride: len(stride) need to be equal to len(inp.shape) |
| 362 | for NCHW, should be (1, 1, stride_h, stride_w) |
| 363 | for NHWC, should be (1, stride_h, stride_w, 1) |
| 364 | kernel: similar to stride, len(kernel) also need to be equal to len(inp.shape) |
| 365 | padding: similar |
| 366 | for NCHW, should be ((0, 0), (0, 0), (pad_h, pad_h), (pad_w, pad_w)) or (0, 0, pad_h, pad_w) |
| 367 | for NHWC, should be ((0, 0), (pad_h, pad_h), (pad_w, pad_w), (0, 0)) or (0, pad_h, pad_w, 0) |
| 368 | """ |
| 369 | ishape, idtype = inp.shape, inp.dtype |
| 370 | assert oshape is not None, "pooling shape infer is not supported" |
| 371 | assert len(ishape) == len(oshape), f"shape error: {ishape} {oshape}" |
| 372 | |
| 373 | def check_param(param, info): |
| 374 | assert len(ishape) == len( |
| 375 | param |
| 376 | ), f"pooling: illegal {info} {param} for {ishape}" |
| 377 | |
| 378 | base_dilation = base_dilation if base_dilation is not None else (1, 1, 1, 1) |
| 379 | kernel_dilation = kernel_dilation if kernel_dilation is not None else (1, 1, 1, 1) |
| 380 | padding = [(p, p) if isinstance(p, int) else p for p in padding] |
| 381 | |
| 382 | check_param(stride, "stride") |
| 383 | check_param(kernel, "kernel") |
| 384 | check_param(padding, "padding") |
| 385 | check_param(base_dilation, "base_dilation") |
| 386 | check_param(kernel_dilation, "kernel_dilation") |
| 387 | |
| 388 | rw = hlo.ReduceWindowOp( |
| 389 | ir_utils.make_ir_type_according_meta_tuple(oshape, idtype), |
| 390 | [inp.tensor], |
| 391 | ir_utils.ir_constant_tuple(unit_factory(idtype)), |
| 392 | ir_utils.dense_int_elements(kernel), |
| 393 | window_strides=ir_utils.dense_int_elements(stride), |
| 394 | base_dilations=ir_utils.dense_int_elements(base_dilation), |
| 395 | window_dilations=ir_utils.dense_int_elements(kernel_dilation), |
| 396 | padding=ir.DenseIntElementsAttr.get( |
| 397 | np.asarray(padding, np.int64), shape=(len(padding), 2) |
| 398 | ), |
| 399 | ) |
| 400 | scalar_type = ir_utils.make_ir_type_according_meta(tuple(), idtype) |
| 401 | reducer_region = rw.regions[0].blocks.append(scalar_type, scalar_type) |
| 402 | with ir.InsertionPoint(reducer_region): |
| 403 | hlo.ReturnOp(reducer(*reducer_region.arguments)) |
| 404 | return HLOTensor(rw.result) |
| 405 |
nothing calls this directly
no test coverage detected