[MLIR][GPU] Add 16-bit version of cudaMemset in cudaRuntimeWrappers

[navdeepkk]

Summary

Add 16-bit version of cudaMemset in cudaRuntimeWrappers and update the GPU to LLVM lowering.

[bondhugula]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

203

Please add in sorted order. One place above.

1376–1380

This logic is wrong and would simply crash. Check if it's int or float first before getting the width.

[navdeepkk]

[MLIR][GPU] Add 16-bit version of cudaMemset in cudaRuntimeWrappers

Add 16-bit version of cudaMemset in cudaRuntimeWrappers and update the GPU to LLVM lowering.

[bondhugula]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1377–1378

Nit: Avoid multiple calls to getIntOr...

[bondhugula]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1396–1398

Use ternary operator for the builder ref: ...& builder = .... ? memset32.... : memset16...; to avoid duplication.

[navdeepkk]

Address comments.

[bondhugula]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1396–1398

This isn't what I meant -- this won't save anything in terms of duplication. Use the ternary operator to assign the right builder.

[navdeepkk]

Resolve comments

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

protected:
   MLIRContext *context = &this->getTypeConverter()->getContext();
 
   Type llvmVoidType = LLVM::LLVMVoidType::get(context);
   LLVM::LLVMPointerType llvmPointerType =
       this->getTypeConverter()->getPointerType(IntegerType::get(context, 8));
   Type llvmPointerPointerType =
       this->getTypeConverter()->getPointerType(llvmPointerType);
   Type llvmInt8Type = IntegerType::get(context, 8);
+  Type llvmInt16Type = IntegerType::get(context, 16);
   Type llvmInt32Type = IntegerType::get(context, 32);
   Type llvmInt64Type = IntegerType::get(context, 64);
   Type llvmInt8PointerType =
       this->getTypeConverter()->getPointerType(llvmInt8Type);
   Type llvmInt64PointerType =
       this->getTypeConverter()->getPointerType(llvmInt64Type);
   Type llvmIntPtrType = IntegerType::get(
       context, this->getTypeConverter()->getPointerBitwidth(0));
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 = {
+  FunctionCallBuilder memset16CallBuilder = {
+      "mgpuMemset16",
+      llvmVoidType,
+      {llvmPointerType /* void *dst */,
+       llvmInt16Type /* unsigned short value */,
+       llvmIntPtrType /* intptr_t sizeBytes */,
+       llvmPointerType /* void *stream */}};
+  FunctionCallBuilder memset32CallBuilder = {
       "mgpuMemset32",
       llvmVoidType,
       {llvmPointerType /* void *dst */, llvmInt32Type /* unsigned int value */,
        llvmIntPtrType /* intptr_t sizeBytes */,
        llvmPointerType /* void *stream */}};
   FunctionCallBuilder setDefaultDeviceCallBuilder = {
       "mgpuSetDefaultDevice",
       llvmVoidType,
       {llvmInt32Type /* uint32_t devIndex */}};
   FunctionCallBuilder createSparseEnvCallBuilder = {
       "mgpuCreateSparseEnv",
       llvmPointerType,
       {llvmPointerType /* void *stream */}};
   FunctionCallBuilder destroySparseEnvCallBuilder = {
LogicalResult ConvertMemsetOpToGpuRuntimeCallPattern::matchAndRewrite(
   if (failed(areAllLLVMTypes(memsetOp, adaptor.getOperands(), rewriter)) ||
       !isConvertibleAndHasIdentityMaps(memRefType) ||
       failed(isAsyncWithOneDependency(rewriter, memsetOp)))
     return failure();
 
   auto loc = memsetOp.getLoc();
 
   Type valueType = adaptor.getValue().getType();
-  if (!valueType.isIntOrFloat() || valueType.getIntOrFloatBitWidth() != 32) {
-    return rewriter.notifyMatchFailure(memsetOp,
-                                       "value must be a 32 bit scalar");
+  unsigned bitWidth = valueType.getIntOrFloatBitWidth();
+  // Ints and floats of 16 or 32 bit width are allowed.
+  if (!valueType.isIntOrFloat() || (bitWidth != 16 && bitWidth != 32)) {
+    return rewriter.notifyMatchFailure(
+        memsetOp, "value must be a 16 or 32 bit int or float");
   }
 
+  unsigned valueTypeWidth = valueType.getIntOrFloatBitWidth();
+  Type bitCastType = valueTypeWidth == 32 ? llvmInt32Type : llvmInt16Type;
+
   MemRefDescriptor dstDesc(adaptor.getDst());
   Value numElements = getNumElements(rewriter, loc, memRefType, dstDesc);
 
   auto value =
-      rewriter.create<LLVM::BitcastOp>(loc, llvmInt32Type, adaptor.getValue());
+      rewriter.create<LLVM::BitcastOp>(loc, bitCastType, adaptor.getValue());
   auto dst = bitAndAddrspaceCast(loc, rewriter, llvmPointerType,
                                  dstDesc.alignedPtr(rewriter, loc),
                                  *getTypeConverter());
 
   auto stream = adaptor.getAsyncDependencies().front();
-  memsetCallBuilder.create(loc, rewriter, {dst, value, numElements, stream});
+  FunctionCallBuilder builder =
+      valueTypeWidth == 32 ? memset32CallBuilder : memset16CallBuilder;
+  builder.create(loc, rewriter, {dst, value, numElements, stream});
 
   rewriter.replaceOp(memsetOp, {stream});
   return success();
 }
 
 LogicalResult ConvertSetDefaultDeviceOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::SetDefaultDeviceOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

 }
 
 extern "C" void mgpuMemset32(void *dst, unsigned int value, size_t count,
                              CUstream stream) {
   CUDA_REPORT_IF_ERROR(cuMemsetD32Async(reinterpret_cast<CUdeviceptr>(dst),
                                         value, count, stream));
 }
 
+extern "C" void mgpuMemset16(void *dst, unsigned short value, size_t count,
+                             CUstream stream) {
+  CUDA_REPORT_IF_ERROR(cuMemsetD16Async(reinterpret_cast<CUdeviceptr>(dst),
+                                        value, count, 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) {

mlir/test/Conversion/GPUCommon/typed-pointers.mlir

func.func @foo(%dst : memref<7xf32, 1>, %src : memref<7xf32>) {
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
 
-  // CHECK: func @foo
-  func.func @foo(%dst : memref<7xf32, 1>, %value : f32) {
+  // CHECK: func @memset_f32
+  func.func @memset_f32(%dst : memref<7xf32, 1>, %value : f32) {
     // CHECK: %[[t0:.*]] = llvm.call @mgpuStreamCreate
     %t0 = gpu.wait async
     // CHECK: %[[size_bytes:.*]] = llvm.mlir.constant
     // CHECK: %[[value:.*]] = llvm.bitcast
     // CHECK: %[[addr_cast:.*]] = llvm.addrspacecast
     // CHECK: %[[dst:.*]] = llvm.bitcast %[[addr_cast]]
     // CHECK: llvm.call @mgpuMemset32(%[[dst]], %[[value]], %[[size_bytes]], %[[t0]])
     %t1 = gpu.memset async [%t0] %dst, %value : memref<7xf32, 1>, f32
     // CHECK: llvm.call @mgpuStreamSynchronize(%[[t0]])
     // CHECK: llvm.call @mgpuStreamDestroy(%[[t0]])
     gpu.wait [%t1]
     return
   }
 }
 
+// -----
+
+module attributes {gpu.container_module} {
+
+  // CHECK: func @memset_f16
+  func.func @memset_f16(%dst : memref<7xf16, 1>, %value : f16) {
+    // CHECK: %[[t0:.*]] = llvm.call @mgpuStreamCreate
+    %t0 = gpu.wait async
+    // CHECK: %[[size_bytes:.*]] = llvm.mlir.constant
+    // CHECK: %[[value:.*]] = llvm.bitcast
+    // CHECK: %[[addr_cast:.*]] = llvm.addrspacecast
+    // CHECK: %[[dst:.*]] = llvm.bitcast %[[addr_cast]]
+    // CHECK: llvm.call @mgpuMemset16(%[[dst]], %[[value]], %[[size_bytes]], %[[t0]])
+    %t1 = gpu.memset async [%t0] %dst, %value : memref<7xf16, 1>, f16
+    // CHECK: llvm.call @mgpuStreamSynchronize(%[[t0]])
+    // CHECK: llvm.call @mgpuStreamDestroy(%[[t0]])
+    gpu.wait [%t1]
+    return
+  }
+}