[mlir] create gpu memset op

[lorenrose1013]

Summary

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

[herhut]

mlir/include/mlir/Dialect/GPU/GPUOps.td

929

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

[lorenrose1013]

mlir/include/mlir/Dialect/GPU/GPUOps.td

929

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]

mlir/include/mlir/Dialect/GPU/GPUOps.td

929

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]

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

[herhut]

Looks great. Just a nit.

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

958

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

[lorenrose1013]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

958

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]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

958

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

[jpienaar]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

958

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

[lorenrose1013]

Verfiy MemsetOp via AllElementsTypeMatch trait.

[lorenrose1013]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

958

Ah thanks, this did the trick!

[herhut]

Thanks!

[lorenrose1013]

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?

mlir/include/mlir/Dialect/GPU/GPUOps.td

def GPU_MemcpyOp : GPU_Op<"memcpy", [GPU_AsyncOpInterface]> {
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $dst`,` $src `:` type($dst)`,` type($src) attr-dict
   }];
   let verifier = [{ return ::verify(*this); }];
   let hasFolder = 1;
 }
 
+def GPU_MemsetOp : GPU_Op<"memset", [GPU_AsyncOpInterface]> {
+
+  let summary = "GPU memset operation";
+
+  let description = [{
+    The `gpu.memset` operation sets the content of memref to a scalar value.
+
+    The op does not execute before all async dependencies have finished
+    executing.
+
+    If the `async` keyword is present, the op is executed asynchronously (i.e.
+    it does not block until the execution has finished on the device). In
+    that case, it returns a !gpu.async.token.
+
+    Example:
+
+    ```mlir
+    %token = gpu.memset async [%dep] %dst, %value : memref<?xf32, 1>, memref<?xf32>
+    ```
+  }];
+
+  let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
+                   Arg<AnyMemRef, "", [MemWrite]>:$dst,
+                   Arg<AnyMemRef, "", [MemRead]>:$value);
+  let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
+
+  let assemblyFormat = [{
+    custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
+    $dst`,` $value `:` type($dst)`,` type($value) attr-dict
+  }];
+  let verifier = [{ return ::verify(*this); }];
+  let hasFolder = 1;
+}
+
 def GPU_SubgroupMmaLoadMatrixOp : GPU_Op<"subgroup_mma_load_matrix",
     [MemoryEffects<[MemRead]>]>{
 
   let summary = "GPU warp synchronous matrix load";
 
   let description = [{
     The `gpu.subgroup_mma_load_matrix` operation loads a matrix collectively
     using all the threads in a subgroup.
let description = [{
 
     The operation takes a scalar input and return a `!gpu.mma_matrix` where each
     element of is equal to the operand constant. The destination mma_matrix type
     must have elememt type equal to the constant type. Since the layout of
     `!gpu.mma_matrix` is opaque this only support setting all the elements to
     the same value.
 
     This op is meant to be used along with `gpu.subgroup_mma_compute`.
-    
     Example:
 
     ```mlir
      %0 = gpu.subgroup_mma_constant_matrix %a :
        !gpu.mma_matrix<16x16xf16, "AOp">
      %1 = gpu.subgroup_mma_constant_matrix %b :
        !gpu.mma_matrix<16x16xf32, "COp">
     ```

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

struct FunctionCallBuilder {
 
   StringRef functionName;
   LLVM::LLVMFunctionType functionType;
 };
 
 template <typename OpTy>
 class ConvertOpToGpuRuntimeCallPattern : public ConvertOpToLLVMPattern<OpTy> {
 public:
-  explicit ConvertOpToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter)
-      : ConvertOpToLLVMPattern<OpTy>(typeConverter) {}
+  explicit ConvertOpToGpuRuntimeCallPattern(LLVMTypeConverter &converter)
+      : ConvertOpToLLVMPattern<OpTy>(converter) {}
 
 protected:
   MLIRContext *context = &this->getTypeConverter()->getContext();
 
   Type llvmVoidType = LLVM::LLVMVoidType::get(context);
   Type llvmPointerType =
       LLVM::LLVMPointerType::get(IntegerType::get(context, 8));
   Type llvmPointerPointerType = LLVM::LLVMPointerType::get(llvmPointerType);
FunctionCallBuilder deallocCallBuilder = {
       llvmVoidType,
       {llvmPointerType /* void *ptr */, llvmPointerType /* void *stream */}};
   FunctionCallBuilder memcpyCallBuilder = {
       "mgpuMemcpy",
       llvmVoidType,
       {llvmPointerType /* void *dst */, llvmPointerType /* void *src */,
        llvmIntPtrType /* intptr_t sizeBytes */,
        llvmPointerType /* void *stream */}};
+  FunctionCallBuilder memsetCallBuilder = {
+      "mgpuMemset",
+      llvmVoidType,
+      {llvmPointerType /* void *dst */, llvmIntPtrType /* void *value */,
+       llvmIntPtrType /* intptr_t sizeBytes */,
+       llvmPointerType /* void *stream */}};
 };
 
 /// A rewrite pattern to convert gpu.host_register operations into a GPU runtime
 /// call. Currently it supports CUDA and ROCm (HIP).
 class ConvertHostRegisterOpToGpuRuntimeCallPattern
     : public ConvertOpToGpuRuntimeCallPattern<gpu::HostRegisterOp> {
 public:
-  ConvertHostRegisterOpToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter)
-      : ConvertOpToGpuRuntimeCallPattern<gpu::HostRegisterOp>(typeConverter) {}
+  ConvertHostRegisterOpToGpuRuntimeCallPattern(LLVMTypeConverter &converter)
+      : ConvertOpToGpuRuntimeCallPattern<gpu::HostRegisterOp>(converter) {}
 
 private:
   LogicalResult
   matchAndRewrite(gpu::HostRegisterOp hostRegisterOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
 /// A rewrite pattern to convert gpu.alloc operations into a GPU runtime
 /// call. Currently it supports CUDA and ROCm (HIP).
 class ConvertAllocOpToGpuRuntimeCallPattern
     : public ConvertOpToGpuRuntimeCallPattern<gpu::AllocOp> {
 public:
-  ConvertAllocOpToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter)
-      : ConvertOpToGpuRuntimeCallPattern<gpu::AllocOp>(typeConverter) {}
+  ConvertAllocOpToGpuRuntimeCallPattern(LLVMTypeConverter &converter)
+      : ConvertOpToGpuRuntimeCallPattern<gpu::AllocOp>(converter) {}
 
 private:
   LogicalResult
   matchAndRewrite(gpu::AllocOp allocOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
 /// A rewrite pattern to convert gpu.dealloc operations into a GPU runtime
 /// call. Currently it supports CUDA and ROCm (HIP).
 class ConvertDeallocOpToGpuRuntimeCallPattern
     : public ConvertOpToGpuRuntimeCallPattern<gpu::DeallocOp> {
 public:
-  ConvertDeallocOpToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter)
-      : ConvertOpToGpuRuntimeCallPattern<gpu::DeallocOp>(typeConverter) {}
+  ConvertDeallocOpToGpuRuntimeCallPattern(LLVMTypeConverter &converter)
+      : ConvertOpToGpuRuntimeCallPattern<gpu::DeallocOp>(converter) {}
 
 private:
   LogicalResult
   matchAndRewrite(gpu::DeallocOp deallocOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
 class ConvertAsyncYieldToGpuRuntimeCallPattern
     : public ConvertOpToGpuRuntimeCallPattern<async::YieldOp> {
 public:
-  ConvertAsyncYieldToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter)
-      : ConvertOpToGpuRuntimeCallPattern<async::YieldOp>(typeConverter) {}
+  ConvertAsyncYieldToGpuRuntimeCallPattern(LLVMTypeConverter &converter)
+      : ConvertOpToGpuRuntimeCallPattern<async::YieldOp>(converter) {}
 
 private:
   LogicalResult
   matchAndRewrite(async::YieldOp yieldOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
 /// A rewrite pattern to convert gpu.wait operations into a GPU runtime
 /// call. Currently it supports CUDA and ROCm (HIP).
 class ConvertWaitOpToGpuRuntimeCallPattern
     : public ConvertOpToGpuRuntimeCallPattern<gpu::WaitOp> {
 public:
-  ConvertWaitOpToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter)
-      : ConvertOpToGpuRuntimeCallPattern<gpu::WaitOp>(typeConverter) {}
+  ConvertWaitOpToGpuRuntimeCallPattern(LLVMTypeConverter &converter)
+      : ConvertOpToGpuRuntimeCallPattern<gpu::WaitOp>(converter) {}
 
 private:
   LogicalResult
   matchAndRewrite(gpu::WaitOp waitOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
 /// A rewrite pattern to convert gpu.wait async operations into a GPU runtime
 /// call. Currently it supports CUDA and ROCm (HIP).
 class ConvertWaitAsyncOpToGpuRuntimeCallPattern
     : public ConvertOpToGpuRuntimeCallPattern<gpu::WaitOp> {
 public:
-  ConvertWaitAsyncOpToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter)
-      : ConvertOpToGpuRuntimeCallPattern<gpu::WaitOp>(typeConverter) {}
+  ConvertWaitAsyncOpToGpuRuntimeCallPattern(LLVMTypeConverter &converter)
+      : ConvertOpToGpuRuntimeCallPattern<gpu::WaitOp>(converter) {}
 
 private:
   LogicalResult
   matchAndRewrite(gpu::WaitOp waitOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
 /// A rewrite patter to convert gpu.launch_func operations into a sequence of
 /// GPU runtime calls. Currently it supports CUDA and ROCm (HIP).
 ///
 /// In essence, a gpu.launch_func operations gets compiled into the following
 /// sequence of runtime calls:
 ///
 /// * moduleLoad        -- loads the module given the cubin / hsaco data
 /// * moduleGetFunction -- gets a handle to the actual kernel function
 /// * getStreamHelper   -- initializes a new compute stream on GPU
 /// * launchKernel      -- launches the kernel on a stream
 /// * streamSynchronize -- waits for operations on the stream to finish
 ///
 /// Intermediate data structures are allocated on the stack.
 class ConvertLaunchFuncOpToGpuRuntimeCallPattern
     : public ConvertOpToGpuRuntimeCallPattern<gpu::LaunchFuncOp> {
 public:
-  ConvertLaunchFuncOpToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter,
+  ConvertLaunchFuncOpToGpuRuntimeCallPattern(LLVMTypeConverter &converter,
                                              StringRef gpuBinaryAnnotation)
-      : ConvertOpToGpuRuntimeCallPattern<gpu::LaunchFuncOp>(typeConverter),
+      : ConvertOpToGpuRuntimeCallPattern<gpu::LaunchFuncOp>(converter),
         gpuBinaryAnnotation(gpuBinaryAnnotation) {}
 
 private:
   Value generateParamsArray(gpu::LaunchFuncOp launchOp,
                             ArrayRef<Value> operands, OpBuilder &builder) const;
   Value generateKernelNameConstant(StringRef moduleName, StringRef name,
                                    Location loc, OpBuilder &builder) const;
 
class EraseGpuModuleOpPattern : public OpRewritePattern<gpu::GPUModuleOp> {
   }
 };
 
 /// A rewrite pattern to convert gpu.memcpy operations into a GPU runtime
 /// call. Currently it supports CUDA and ROCm (HIP).
 class ConvertMemcpyOpToGpuRuntimeCallPattern
     : public ConvertOpToGpuRuntimeCallPattern<gpu::MemcpyOp> {
 public:
-  ConvertMemcpyOpToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter)
-      : ConvertOpToGpuRuntimeCallPattern<gpu::MemcpyOp>(typeConverter) {}
+  ConvertMemcpyOpToGpuRuntimeCallPattern(LLVMTypeConverter &converter)
+      : ConvertOpToGpuRuntimeCallPattern<gpu::MemcpyOp>(converter) {}
 
 private:
   LogicalResult
   matchAndRewrite(gpu::MemcpyOp memcpyOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
+
+/// A rewrite pattern to convert gpu.memset operations into a GPU runtime
+/// call. Currently it supports CUDA and ROCm (HIP).
+class ConvertMemsetOpToGpuRuntimeCallPattern
+    : public ConvertOpToGpuRuntimeCallPattern<gpu::MemsetOp> {
+public:
+  ConvertMemsetOpToGpuRuntimeCallPattern(LLVMTypeConverter &converter)
+      : ConvertOpToGpuRuntimeCallPattern<gpu::MemsetOp>(converter) {}
+
+private:
+  LogicalResult
+  matchAndRewrite(gpu::MemsetOp memsetOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override;
+};
 } // namespace
 
 void GpuToLLVMConversionPass::runOnOperation() {
   LLVMTypeConverter converter(&getContext());
   RewritePatternSet patterns(&getContext());
   LLVMConversionTarget target(getContext());
 
   target.addIllegalDialect<gpu::GPUDialect>();
   target.addIllegalOp<UnrealizedConversionCastOp>();
 
   populateVectorToLLVMConversionPatterns(converter, patterns);
   populateMemRefToLLVMConversionPatterns(converter, patterns);
   populateStdToLLVMConversionPatterns(converter, patterns);
   populateAsyncStructuralTypeConversionsAndLegality(converter, patterns,
                                                     target);
 
   converter.addConversion(
       [context = &converter.getContext()](gpu::AsyncTokenType type) -> Type {
         return LLVM::LLVMPointerType::get(IntegerType::get(context, 8));
       });
   patterns.add<ConvertAllocOpToGpuRuntimeCallPattern,
                ConvertDeallocOpToGpuRuntimeCallPattern,
                ConvertHostRegisterOpToGpuRuntimeCallPattern,
                ConvertMemcpyOpToGpuRuntimeCallPattern,
+               ConvertMemsetOpToGpuRuntimeCallPattern,
                ConvertWaitAsyncOpToGpuRuntimeCallPattern,
                ConvertWaitOpToGpuRuntimeCallPattern,
                ConvertAsyncYieldToGpuRuntimeCallPattern>(converter);
   patterns.add<ConvertLaunchFuncOpToGpuRuntimeCallPattern>(converter,
                                                            gpuBinaryAnnotation);
   patterns.add<EraseGpuModuleOpPattern>(&converter.getContext());
 
   if (failed(
LogicalResult ConvertMemcpyOpToGpuRuntimeCallPattern::matchAndRewrite(
   auto stream = adaptor.asyncDependencies().front();
   memcpyCallBuilder.create(loc, rewriter, {dst, src, sizeBytes, stream});
 
   rewriter.replaceOp(memcpyOp, {stream});
 
   return success();
 }
 
+LogicalResult ConvertMemsetOpToGpuRuntimeCallPattern::matchAndRewrite(
+    gpu::MemsetOp memsetOp, ArrayRef<Value> operands,
+    ConversionPatternRewriter &rewriter) const {
+  auto memRefType = memsetOp.dst().getType().cast<MemRefType>();
+
+  if (failed(areAllLLVMTypes(memsetOp, operands, rewriter)) ||
+      !isConvertibleAndHasIdentityMaps(memRefType) ||
+      failed(isAsyncWithOneDependency(rewriter, memsetOp)))
+    return failure();
+
+  auto loc = memsetOp.getLoc();
+  auto adaptor = gpu::MemsetOpAdaptor(operands, memsetOp->getAttrDictionary());
+
+  MemRefDescriptor dstDesc(adaptor.dst());
+
+  Value numElements =
+      memRefType.hasStaticShape()
+          ? createIndexConstant(rewriter, loc, memRefType.getNumElements())
+          // For identity layouts (verified above), the number of elements is
+          // stride[0] * size[0].
+          : rewriter.create<LLVM::MulOp>(loc, dstDesc.stride(rewriter, loc, 0),
+                                         dstDesc.size(rewriter, loc, 0));
+
+  Type elementPtrType = getElementPtrType(memRefType);
+  Value nullPtr = rewriter.create<LLVM::NullOp>(loc, elementPtrType);
+  Value sizeBytesGepPtr = rewriter.create<LLVM::GEPOp>(
+      loc, elementPtrType, ArrayRef<Value>{nullPtr, numElements});
+  Value valueGepPtr = rewriter.create<LLVM::GEPOp>(
+      loc, elementPtrType,
+      ArrayRef<Value>{
+          MemRefDescriptor(adaptor.value()).alignedPtr(rewriter, loc),
+          createIndexConstant(rewriter, loc, 0)});
+
+  auto sizeBytes =
+      rewriter.create<LLVM::PtrToIntOp>(loc, getIndexType(), sizeBytesGepPtr);
+  auto value =
+      rewriter.create<LLVM::PtrToIntOp>(loc, getIndexType(), valueGepPtr);
+  auto dst = rewriter.create<LLVM::BitcastOp>(
+      loc, llvmPointerType, dstDesc.alignedPtr(rewriter, loc));
+
+  auto stream = adaptor.asyncDependencies().front();
+  memsetCallBuilder.create(loc, rewriter, {dst, value, sizeBytes, stream});
+
+  rewriter.replaceOp(memsetOp, {stream});
+  return success();
+}
+
 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
 mlir::createGpuToLLVMConversionPass() {
   return std::make_unique<GpuToLLVMConversionPass>();
 }

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

if (getElementTypeOrSelf(srcType) != getElementTypeOrSelf(dstType))
     return op.emitOpError("arguments have incompatible element type");
 
   if (failed(verifyCompatibleShape(srcType, dstType)))
     return op.emitOpError("arguments have incompatible shape");
 
   return success();
 }
 
+static LogicalResult verify(MemsetOp op) {
+  auto dstType = op.dst().getType();
+  auto valueType = op.value().getType();
+
+  if (getElementTypeOrSelf(valueType) != getElementTypeOrSelf(dstType))
+    return op.emitOpError("arguments have incompatible element type");
+
+  ArrayRef<int64_t> value_shape = valueType.dyn_cast<ShapedType>().getShape();
+  if (value_shape.size() != 1 || value_shape.front() != 1)
+    return op.emitOpError("value argument must be a scalar");
+
+  return success();
+}
+
 static ParseResult parseAsyncDependencies(
     OpAsmParser &parser, Type &asyncTokenType,
     SmallVectorImpl<OpAsmParser::OperandType> &asyncDependencies) {
   auto loc = parser.getCurrentLocation();
   if (succeeded(parser.parseOptionalKeyword("async"))) {
     if (parser.getNumResults() == 0)
       return parser.emitError(loc, "needs to be named when marked 'async'");
     asyncTokenType = parser.getBuilder().getType<AsyncTokenType>();
static LogicalResult foldMemRefCast(Operation *op) {
   return success(folded);
 }
 
 LogicalResult MemcpyOp::fold(ArrayRef<Attribute> operands,
                              SmallVectorImpl<::mlir::OpFoldResult> &results) {
   return foldMemRefCast(*this);
 }
 
+LogicalResult MemsetOp::fold(ArrayRef<Attribute> operands,
+                             SmallVectorImpl<::mlir::OpFoldResult> &results) {
+  return foldMemRefCast(*this);
+}
+
 #include "mlir/Dialect/GPU/GPUOpInterfaces.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/GPU/GPUOps.cpp.inc"

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

 
 #include <cassert>
 #include <numeric>
 
 #include "mlir/ExecutionEngine/CRunnerUtils.h"
 #include "llvm/ADT/ArrayRef.h"
 
 #include "cuda.h"
+#include "cuda_runtime_api.h"
 
 #ifdef _WIN32
 #define MLIR_CUDA_WRAPPERS_EXPORT __declspec(dllexport)
 #else
 #define MLIR_CUDA_WRAPPERS_EXPORT
 #endif // _WIN32
 
 #define CUDA_REPORT_IF_ERROR(expr)                                             \
   [](CUresult result) {                                                        \
     if (!result)                                                               \
       return;                                                                  \
     const char *name = nullptr;                                                \
     cuGetErrorName(result, &name);                                             \
     if (!name)                                                                 \
       name = "<unknown>";                                                      \
     fprintf(stderr, "'%s' failed with '%s'\n", #expr, name);                   \
   }(expr)
 
+#define CUDA_RUNTIME_REPORT_IF_ERROR(expr)                                     \
+  [](cudaError_t result) {                                                     \
+    if (result == cudaSuccess)                                                 \
+      return;                                                                  \
+    const char *name = cudaGetErrorName(result);                               \
+    if (!name)                                                                 \
+      name = "<unknown>";                                                      \
+    fprintf(stderr, "'%s' failed with '%s'\n", #expr, name);                   \
+  }(expr)
+
 // Make the primary context of device 0 current for the duration of the instance
 // and restore the previous context on destruction.
 class ScopedContext {
 public:
   ScopedContext() {
     // Static reference to CUDA primary context for device ordinal 0.
     static CUcontext context = [] {
       CUDA_REPORT_IF_ERROR(cuInit(/*flags=*/0));

 
 extern "C" void mgpuMemcpy(void *dst, void *src, uint64_t sizeBytes,
                            CUstream stream) {
   CUDA_REPORT_IF_ERROR(cuMemcpyAsync(reinterpret_cast<CUdeviceptr>(dst),
                                      reinterpret_cast<CUdeviceptr>(src),
                                      sizeBytes, stream));
 }
 
+extern "C" void mgpuMemset(void *dst, int value, uint64_t sizeBytes,
+                           CUstream stream) {
+  CUDA_RUNTIME_REPORT_IF_ERROR(cudaMemsetAsync(dst, value, sizeBytes, stream));
+}
+
 /// Helper functions for writing mlir example code
 
 // Allows to register byte array with the CUDA runtime. Helpful until we have
 // transfer functions implemented.
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
 mgpuMemHostRegister(void *ptr, uint64_t sizeBytes) {
   ScopedContext scopedContext;
   CUDA_REPORT_IF_ERROR(cuMemHostRegister(ptr, sizeBytes, /*flags=*/0));

mlir/lib/ExecutionEngine/RocmRuntimeWrappers.cpp

 //===----------------------------------------------------------------------===//
 
 #include <cassert>
 #include <numeric>
 
 #include "mlir/ExecutionEngine/CRunnerUtils.h"
 #include "llvm/ADT/ArrayRef.h"
 
 #include "hip/hip_runtime.h"
 
 #define HIP_REPORT_IF_ERROR(expr)                                              \
   [](hipError_t result) {                                                      \
     if (!result)                                                               \
       return;                                                                  \
     const char *name = hipGetErrorName(result);                                \
     if (!name)                                                                 \
       name = "<unknown>";                                                      \

 }
 
 extern "C" void mgpuMemcpy(void *dst, void *src, uint64_t sizeBytes,
                            hipStream_t stream) {
   HIP_REPORT_IF_ERROR(
       hipMemcpyAsync(dst, src, sizeBytes, hipMemcpyDefault, stream));
 }
 
+extern "C" void mgpuMemset(void *dst, int value, uint64_t sizeBytes,
+                           uint64_t width, hipStream_t stream) {
+  HIP_REPORT_IF_ERROR(hipMemsetAsync(dst, value, sizeBytes, stream));
+}
 /// Helper functions for writing mlir example code
 
 // Allows to register byte array with the ROCM runtime. Helpful until we have
 // transfer functions implemented.
 extern "C" void mgpuMemHostRegister(void *ptr, uint64_t sizeBytes) {
   ScopedContext scopedContext;
   HIP_REPORT_IF_ERROR(hipHostRegister(ptr, sizeBytes, /*flags=*/0));
 }

mlir/test/Conversion/GPUCommon/lower-memset-to-gpu-runtime-calls.mlir

+// RUN: mlir-opt %s --gpu-to-llvm | FileCheck %s
+
+module attributes {gpu.container_module} {
+
+  // CHECK: func @foo
+  func @foo(%dst : memref<7xf32, 1>, %value : memref<1xf32>) {
+    // CHECK: %[[t0:.*]] = llvm.call @mgpuStreamCreate
+    %t0 = gpu.wait async
+    // CHECK: %[[size_bytes:.*]] = llvm.ptrtoint
+    // CHECK: %[[value:.*]] = llvm.ptrtoint
+    // CHECK: %[[dst:.*]] = llvm.bitcast
+    // CHECK: llvm.call @mgpuMemset(%[[dst]], %[[value]], %[[size_bytes]], %[[t0]])
+    %t1 = gpu.memset async [%t0] %dst, %value : memref<7xf32, 1>, memref<1xf32>
+    // CHECK: llvm.call @mgpuStreamSynchronize(%[[t0]])
+    // CHECK: llvm.call @mgpuStreamDestroy(%[[t0]])
+    gpu.wait [%t1]
+    return
+  }
+}

mlir/test/Dialect/GPU/canonicalize.mlir

 // RUN: mlir-opt %s -canonicalize --split-input-file -allow-unregistered-dialect | FileCheck %s
 
 // CHECK-LABEL: @memcpy_after_cast
 func @memcpy_after_cast(%arg0: memref<10xf32>, %arg1: memref<10xf32>) {
   // CHECK-NOT: memref.cast
   // CHECK: gpu.memcpy
   %0 = memref.cast %arg0 : memref<10xf32> to memref<?xf32>
   %1 = memref.cast %arg1 : memref<10xf32> to memref<?xf32>
   gpu.memcpy %0,%1 : memref<?xf32>, memref<?xf32>
   return
 }
+
+// CHECK-LABEL: @memset_after_cast
+func @memset_after_cast(%arg0: memref<10xf32>, %arg1: memref<1xf32>) {
+  // CHECK-NOT: memref.cast
+  // CHECK: gpu.memset
+  %0 = memref.cast %arg0 : memref<10xf32> to memref<?xf32>
+  %1 = memref.cast %arg1 : memref<1xf32> to memref<1xf32>
+  gpu.memset %0,%1 : memref<?xf32>, memref<1xf32>
+  return
+}

mlir/test/Dialect/GPU/invalid.mlir

 
 func @memcpy_incompatible_shape(%dst : memref<7xf32>, %src : memref<9xf32>) {
   // expected-error @+1 {{'gpu.memcpy' op arguments have incompatible shape}}
   gpu.memcpy %dst, %src  : memref<7xf32>, memref<9xf32>
 }
 
 // -----
 
+func @memset_incompatible_type(%dst : memref<?xf32>, %value : memref<1xi32>) {
+  // expected-error @+1 {{'gpu.memset' op arguments have incompatible element type}}
+  gpu.memset %dst, %value  : memref<?xf32>, memref<1xi32>
+}
+
+// -----
+
+func @memset_incompatible_shape(%dst : memref<?xf32>, %value : memref<9xf32>) {
+  // expected-error @+1 {{'gpu.memset' op value argument must be a scalar}}
+  gpu.memset %dst, %value  : memref<?xf32>, memref<9xf32>
+}
+
+// -----
+
 func @mmamatrix_invalid_shape(){
     %wg = memref.alloca() {alignment = 32} : memref<32x32xf16, 3>
     %i = constant 16 : index
     // expected-error @+1 {{MMAMatrixType must have exactly two dimensions}}
     %0 = gpu.subgroup_mma_load_matrix %wg[%i, %i] {leadDimension = 32 : index} : memref<32x32xf16, 3> -> !gpu.mma_matrix<16x16x16xf16, "AOp">
     return
 }
 

mlir/test/Dialect/GPU/ops.mlir

func @memcpy(%dst : memref<3x7xf32>, %src : memref<3x7xf32, 1>) {
     gpu.memcpy %dst, %src : memref<3x7xf32>, memref<3x7xf32, 1>
     // CHECK: %[[t0:.*]] = gpu.wait async
     %0 = gpu.wait async
     // CHECK: {{.*}} = gpu.memcpy async [%[[t0]]] {{.*}}, {{.*}} : memref<3x7xf32>, memref<3x7xf32, 1>
     %1 = gpu.memcpy async [%0] %dst, %src : memref<3x7xf32>, memref<3x7xf32, 1>
     return
   }
 
+  func @memset(%dst : memref<3x7xf32>, %value : memref<1xf32, 1>) {
+    // CHECK-LABEL: func @memset
+    // CHECK: gpu.memset {{.*}}, {{.*}} : memref<3x7xf32>, memref<1xf32, 1>
+    gpu.memset %dst, %value : memref<3x7xf32>, memref<1xf32, 1>
+    // CHECK: %[[t0:.*]] = gpu.wait async
+    %0 = gpu.wait async
+    // CHECK: {{.*}} = gpu.memset async [%[[t0]]] {{.*}}, {{.*}} : memref<3x7xf32>, memref<1xf32, 1>
+    %1 = gpu.memset async [%0] %dst, %value : memref<3x7xf32>, memref<1xf32, 1>
+    return
+  }
+
   func @mmamatrix_valid_element_type(){
     // CHECK-LABEL: func @mmamatrix_valid_element_type
     %wg = memref.alloca() {alignment = 32} : memref<32x32xf16, 3>
     // CHECK: %[[wg:.*]] = memref.alloca()
     %i = constant 16 : index
     // CHECK: %[[i:.*]] = constant 16 : index
      %cst = constant 1.000000e+00 : f32
     // CHECK: %[[cst:.*]] = constant 1.000000e+00 : f32
     %0 = gpu.subgroup_mma_load_matrix %wg[%i, %i] {leadDimension = 32 : index} : memref<32x32xf16, 3> -> !gpu.mma_matrix<16x16xf16, "AOp">
     // CHECK: gpu.subgroup_mma_load_matrix %[[wg]][%[[i]], %[[i]]] {leadDimension = 32 : index} : memref<32x32xf16, 3> -> !gpu.mma_matrix<16x16xf16, "AOp">
     %1 = gpu.subgroup_mma_constant_matrix %cst : !gpu.mma_matrix<16x16xf32, "COp">
     // CHECK: gpu.subgroup_mma_constant_matrix %[[cst]] : !gpu.mma_matrix<16x16xf32, "COp">
     return
   }
 }