Page MenuHomePhabricator

[mlir][Linalg] Add 3D pooling named ops to Linalg.
ClosedPublic

Authored by hanchung on Jul 1 2021, 7:02 PM.

Diff Detail

Event Timeline

hanchung created this revision.Jul 1 2021, 7:02 PM
hanchung requested review of this revision.Jul 1 2021, 7:02 PM
Herald added a project: Restricted Project. · View Herald TranscriptJul 1 2021, 7:02 PM

I found that there is a pooling_nhwc_sum_poly in LinalgNamedStructuredOps.yaml, but not pooling_nhwc_min/max_poly. I'm not sure if pooling min/max ops can be added to yaml file, so I still add 3D poolings to the tc file.

Let me know if you'd like me to add them to yaml file. I'm happy to learn it. :D

gysit added a comment.Jul 2 2021, 12:53 AM

I found that there is a pooling_nhwc_sum_poly in LinalgNamedStructuredOps.yaml, but not pooling_nhwc_min/max_poly. I'm not sure if pooling min/max ops can be added to yaml file, so I still add 3D poolings to the tc file.

Let me know if you'd like me to add them to yaml file. I'm happy to learn it. :D

I just landed https://reviews.llvm.org/D105203 which contains max pooling operation but did not yet add support for the min operation. I hope I can come up with a revision implementing the min operator today.

Also note that pooling_nhwc_sum_poly are test operations I added. Feel free to extend modify them so that they fit rest. E.g. you may want another naming scheme and replace poly with something else.

hanchung updated this revision to Diff 357595.Jul 9 2021, 12:25 PM

Add ops to yaml file instead

hanchung accepted this revision.Jul 9 2021, 12:27 PM

Thanks @gysit, I added them to yaml files instead, PTAL!

I will migrate the 2D pooling ops to yaml files and replace all the uses in MHLO/IREE repo in a later patch.

This revision is now accepted and ready to land.Jul 9 2021, 12:27 PM
hanchung requested review of this revision.Jul 9 2021, 12:29 PM
gysit added a comment.Jul 12 2021, 4:12 AM

Thanks @gysit, I added them to yaml files instead, PTAL!

I will migrate the 2D pooling ops to yaml files and replace all the uses in MHLO/IREE repo in a later patch.

Great!

Did you also edit llvm-project/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py? The YAML file is machine generated and shouldn't be edited directly. Instead one can edit core_named_ops.py and follow the procedure discussed by the basic usage section of https://mlir.llvm.org/docs/Dialects/Linalg/OpDSL/.

The OpDSL code to generate the 3D pooling operations should look as follows:

@linalg_structured_op
def pooling_ndhwc_sum(
    I=TensorDef(T1, S.N, S.D, S.H, S.W, S.C),
    K=TensorDef(T2, S.KD, S.KH, S.KW, index_dims=[D.kd, D.kh, D.kw]),
    O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, output=True),
    strides=AttributeDef(S.SD, S.SH, S.SW),
    dilations=AttributeDef(S.DD, S.DH, S.DW)):
  """Performs 3D sum pooling.

  Numeric casting is performed on the input operand, promoting it to the same
  data type as the accumulator/output.
  """
  domain(D.n, D.od, D.oh, D.ow, D.kd, D.kh, D.kw, D.c)
  O[D.n, D.od, D.oh, D.ow, D.c] += cast(
          U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH,
               D.ow * S.SW + D.kw * S.DW, D.c])


@linalg_structured_op
def pooling_ndhwc_max(
    I=TensorDef(T1, S.N, S.D, S.H, S.W, S.C),
    K=TensorDef(T2, S.KD, S.KH, S.KW, index_dims=[D.kd, D.kh, D.kw]),
    O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, output=True),
    strides=AttributeDef(S.SD, S.SH, S.SW),
    dilations=AttributeDef(S.DD, S.DH, S.DW)):
  """Performs 3D max pooling.

  Numeric casting is performed on the input operand, promoting it to the same
  data type as the accumulator/output.
  """
  domain(D.n, D.od, D.oh, D.ow, D.kd, D.kh, D.kw, D.c)
  O[D.n, D.od, D.oh, D.ow, D.c] = ReduceFn.max(D.kd, D.kh, D.kw)(
      cast(
          U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH,
               D.ow * S.SW + D.kw * S.DW, D.c]))


@linalg_structured_op
def pooling_ndhwc_min(
    I=TensorDef(T1, S.N, S.D, S.H, S.W, S.C),
    K=TensorDef(T2, S.KD, S.KH, S.KW, index_dims=[D.kd, D.kh, D.kw]),
    O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, output=True),
    strides=AttributeDef(S.SD, S.SH, S.SW),
    dilations=AttributeDef(S.DD, S.DH, S.DW)):
  """Performs 3D max pooling.

  Numeric casting is performed on the input operand, promoting it to the same
  data type as the accumulator/output.
  """
  domain(D.n, D.od, D.oh, D.ow, D.kd, D.kh, D.kw, D.c)
  O[D.n, D.od, D.oh, D.ow, D.c] = ReduceFn.min(D.kd, D.kh, D.kw)(
      cast(
          U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH,
               D.ow * S.SW + D.kw * S.DW, D.c]))
gysit added a comment.Jul 12 2021, 4:19 AM

PS there are already some 2d pooling op examples in core_named_ops.py. I currently added the postfix _poly to avoid a name collisions with the existing ops. I think once the existing ops are removed we can just drop the _poly postfix? I also added some initial integration tests for these ops. I don't think we should test all ops there but I just wanted to make you aware of these tests in case you want to verify things:

end-to-end integration test:
llvm-project/mlir/test/python/integration/dialects/linalg/opsrun.py

tests to verify the code generated from the yaml specification:
llvm-project/mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir

hanchung updated this revision to Diff 357995.Jul 12 2021, 11:04 AM

Follow commnets to generate the yaml file.

Thanks! I added the defs to core_named_ops.py and copy-paste the result the yaml file! Also moved the mlir roundtrip tests from name-ops to the poly ops file.

hanchung updated this revision to Diff 358025.Jul 12 2021, 12:20 PM

Move roundtrip tests back to named-ops.mlir

gysit accepted this revision.Jul 12 2021, 12:23 PM

Thanks! LGTM

This revision is now accepted and ready to land.Jul 12 2021, 12:23 PM
This revision was automatically updated to reflect the committed changes.