Page MenuHomePhabricator

[mlir] create gpu memset op

Authored by lorenrose1013 on Aug 5 2021, 3:01 AM.



Create a gpu memset op and corresponding CUDA and ROCm wrappers.

Diff Detail

Event Timeline

lorenrose1013 created this revision.Aug 5 2021, 3:01 AM
lorenrose1013 requested review of this revision.Aug 5 2021, 3:01 AM
herhut added inline comments.Aug 5 2021, 5:14 AM

If this has to be a scalar, why is it not a scalar value but a memref?

lorenrose1013 added inline comments.Aug 10 2021, 2:06 AM

This is my first time working with LLVM/MLIR, and I based this off of the Memcpy operation above. Would Arg<AnyTypeOf<[F16, F32, I32], be a better replacement, or is there a more appropriate way to type the scalar?

csigg added inline comments.Aug 10 2021, 3:13 AM

I would probably infer the type of $value from the the element type of $dst (custom<GpuMemsetType>(type($dst), type($value)) in the assembly format, and provide print/parseGpuMemSetType()).

Alternatively, use AnyType for $value and check that it matches the element type of $dst in the verifier.

lorenrose1013 updated this revision to Diff 368050.EditedAug 23 2021, 2:35 AM

Updating D107548: [mlir] create gpu memset op

Made changes based on internal review:

  • explicitly only supporting 32 bit memset for now
  • cuda wrappers call cuMemsetD32Async directly and no longer import the cuda runtime api
  • extracted out getNumElements for gpu memrefs for logic that was common to memset and memcpy
  • mlir memset op takes in a scalar of AnyType instead of AnyMemRef

Looks great. Just a nit.


You could also use the SameOperandsElementType trait in the op definition to have this check autogenerated,

lorenrose1013 added inline comments.Aug 24 2021, 4:18 AM

I've tried this but the tests are failing now, I think because the first operand is actually the async token and so it doesn't match the dst and value types?

jpienaar added inline comments.

There is a parametrized one (on mobile, will link once in) where one can specify beyond all matching.

jpienaar added inline comments.Aug 24 2021, 8:05 AM

AllElementTypesMatch<["dst", "value"]> - you are looking at element type of one and type of another, but I think this should do that.

Verfiy MemsetOp via AllElementsTypeMatch trait.

lorenrose1013 added inline comments.Aug 24 2021, 8:44 AM

Ah thanks, this did the trick!

herhut accepted this revision.Aug 26 2021, 2:34 AM


This revision is now accepted and ready to land.Aug 26 2021, 2:34 AM
lorenrose1013 accepted this revision.Aug 26 2021, 3:38 AM

Thanks for the approval!

I've tried to land this, but I don't seem to have github permissions for llvm (Permission to llvm/llvm-project.git denied to lorenrose1013.)
Is this something I can request, or is someone else able to land the change instead?

This revision was landed with ongoing or failed builds.Sep 3 2021, 11:13 PM
Closed by commit rG361458b1ce89: [mlir] create gpu memset op (authored by lorenrose1013, committed by csigg). · Explain Why
This revision was automatically updated to reflect the committed changes.