[mlir][NVGPU] Adding Support for cp_async_zfill via Inline Asm

[manishucsd]

Summary

cp_async_zfill is currently not present in the nvvm backend, this patch adds cp_async_zfill support by adding inline asm when lowering from nvgpu to nvvm.

[manishucsd]

Minor update

[manishucsd]

clang format

[ThomasRaoux]

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

388

nit: remove empty line

417

I believe this naming is cutlass specific and I don't really see it anywhere in ptx or nvvm spec. I think it would be better to explain the difference between those two different version of cp.async and not name it cp_async_zfill

[christopherbate]

Great thanks for adding this!

Just a small question: What is the reason for using inline ASM? If you haven't already tried, adding the NVPTX backend support for intrinsics is pretty trivial. It's probably the same amount of code. If you search for past diffs that add NVPTX intrinsics, you can look at the files they modify and follow the same process.

[manishucsd]

Review comments.

[manishucsd]

[ThomasRaoux]

Looks great!

[ThomasRaoux]

In D132269#3740155, @christopherbate wrote:

Great thanks for adding this!

Just a small question: What is the reason for using inline ASM? If you haven't already tried, adding the NVPTX backend support for intrinsics is pretty trivial. It's probably the same amount of code. If you search for past diffs that add NVPTX intrinsics, you can look at the files they modify and follow the same process.

I had suggested it that way to avoid having to add an intrinsic at this point as it would require adding a bunch of new intrinsics (there are 5 for cp.aync that would double it). I wonder if we want to duplicate the intrinsics or let the codegen match cp.async with cp_size == src_size. Anyway I was trying to avoid this discussion and effort at the moment. Maybe I overestimated the effort needed.

In any case I agree that we should add support to it in NVPTX long term.

[manishucsd]

In D132269#3740155, @christopherbate wrote:

Great thanks for adding this!

Just a small question: What is the reason for using inline ASM? If you haven't already tried, adding the NVPTX backend support for intrinsics is pretty trivial. It's probably the same amount of code. If you search for past diffs that add NVPTX intrinsics, you can look at the files they modify and follow the same process.

For supporting async copy with zero fills, inline asm seems to be the fastest route for now. Yeah, we would want to eventually support it via NVPTX backend.

[manishucsd]

Fixes discovered as I try and use this patch in codegen

mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td

let description = [{
     This is memory access will be pending to be added to a group.
 
     This op is meant to be used with `gpu.device_async_create_group` and
     `gpu.device_async_wait` to synchronize copies as explained in those ops
     descriptions. 
     `bypassL1` attribute is hint to the backend and hardware that
     the copy should by pass the L1 cache, this may be dropped by the backend or
     hardware. 
+    `dstElements` attribute is the total number of elements written to 
+    destination (shared memory).
+    `srcElements` argument is the total number of elements read from 
+    source (global memory).
+    
+    srcElements` is an optional argument and when present it only reads 
+    srcElements number of elements from the source global memory and zero fills 
+    the rest of the elements in the destination shared memory.
 
     In order to do a copy and wait for the result we need the following
     combination:
     ```
     // copy 1.
     %cp1 = gpu.device_async_copy %A[%c0], %B[%c0], 4 :memref<16xf32> to memref<16xf32, 3>
     // copy 2.
     %cp2 = gpu.device_async_copy %C[%c0], %D[%c0], 4 : memref<16xf32> to memref<16xf32, 3>
%0 = gpu.device_async_copy %src[%c0, %c0], %dst[%c0, %c0, %c0], 4 :
       memref<4x5xf32> to memref<2x7x5xf32, 3>
     ```
   }];
   let results = (outs NVGPU_DeviceAsyncToken:$asyncToken);
   let arguments = (ins Arg<AnyMemRef, "", [MemWrite]>:$dst,
                        Variadic<Index>:$dstIndices,
                        Arg<AnyMemRef, "", [MemRead]>:$src,
                        Variadic<Index>:$srcIndices,
-                       IndexAttr:$numElements,
+                       IndexAttr:$dstElements,
+                       Optional<Index>:$srcElements,
                        OptionalAttr<UnitAttr>:$bypassL1);
   let assemblyFormat = [{
-    $src `[` $srcIndices `]` `,` $dst `[` $dstIndices `]` `,` $numElements
+    $src `[` $srcIndices `]` `,` $dst `[` $dstIndices `]` `,` $dstElements (`,` $srcElements^)?
       attr-dict `:` type($src) `to` type($dst)
   }];
   let hasVerifier = 1;
 }
 
 def NVGPU_DeviceAsyncCreateGroupOp : NVGPU_Op<"device_async_create_group", []> {
   let summary = "device side asynchronous create group operation";
   let description = [{

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

void runOnOperation() override {
     target.addLegalDialect<::mlir::LLVM::LLVMDialect>();
     target.addLegalDialect<::mlir::NVVM::NVVMDialect>();
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };
 
+static void emitCpAsyncOpZfillAsm(Location loc, Value dstPtr, Value srcPtr,
+                                  Value dstBytes, Value srcElements,
+                                  mlir::MemRefType elementType,
+                                  ConversionPatternRewriter &rewriter) {
+  auto asmDialectAttr = LLVM::AsmDialectAttr::get(rewriter.getContext(),
+                                                  LLVM::AsmDialect::AD_ATT);
+  const char *asmStr = "cp.async.cg.shared.global [$0], [$1], $2, $3;\n";
+  const char *asmConstraints = "r,l,n,r";
+
+  Value c3I32 = rewriter.create<LLVM::ConstantOp>(
+      loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(3));
+  Value bitwidth = rewriter.create<LLVM::ConstantOp>(
+      loc, rewriter.getI32Type(),
+      rewriter.getI32IntegerAttr(elementType.getElementTypeBitWidth()));
+  Value srcElementsI32 =
+      rewriter.create<LLVM::TruncOp>(loc, rewriter.getI32Type(), srcElements);
+  Value srcBytes = rewriter.create<LLVM::LShrOp>(
+      loc, rewriter.create<LLVM::MulOp>(loc, bitwidth, srcElementsI32), c3I32);
+
+  SmallVector<Value> asmVals{dstPtr, srcPtr, dstBytes, srcBytes};
+
+  rewriter.create<LLVM::InlineAsmOp>(
+      loc, LLVM::LLVMVoidType::get(rewriter.getContext()), /*operands=*/asmVals,
+      /*asm_string=*/asmStr,
+      /*constraints=*/asmConstraints, /*has_side_effects=*/true,
+      /*is_align_stack=*/false, /*asm_dialect=*/asmDialectAttr,
+      /*operand_attrs=*/ArrayAttr());
+}
+
 struct NVGPUAsyncCopyLowering
     : public ConvertOpToLLVMPattern<nvgpu::DeviceAsyncCopyOp> {
   using ConvertOpToLLVMPattern<
       nvgpu::DeviceAsyncCopyOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(nvgpu::DeviceAsyncCopyOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
     auto dstMemrefType = op.getDst().getType().cast<MemRefType>();
     Value dstPtr = getStridedElementPtr(loc, dstMemrefType, adaptor.getDst(),
                                         adaptor.getDstIndices(), rewriter);
     auto i8Ty = IntegerType::get(op.getContext(), 8);
     auto dstPointerType =
         LLVM::LLVMPointerType::get(i8Ty, dstMemrefType.getMemorySpaceAsInt());
     dstPtr = rewriter.create<LLVM::BitcastOp>(loc, dstPointerType, dstPtr);
 
     auto srcMemrefType = op.getSrc().getType().cast<MemRefType>();
 
     Value scrPtr = getStridedElementPtr(loc, srcMemrefType, adaptor.getSrc(),
                                         adaptor.getSrcIndices(), rewriter);
     auto srcPointerType =
         LLVM::LLVMPointerType::get(i8Ty, srcMemrefType.getMemorySpaceAsInt());
     scrPtr = rewriter.create<LLVM::BitcastOp>(loc, srcPointerType, scrPtr);
     // Intrinsics takes a global pointer so we need an address space cast.
     auto srcPointerGlobalType = LLVM::LLVMPointerType::get(
         i8Ty, NVVM::NVVMMemorySpace::kGlobalMemorySpace);
     scrPtr = rewriter.create<LLVM::AddrSpaceCastOp>(loc, srcPointerGlobalType,
                                                     scrPtr);
-    int64_t numElements = adaptor.getNumElements().getZExtValue();
+    int64_t dstElements = adaptor.getDstElements().getZExtValue();
     int64_t sizeInBytes =
-        (dstMemrefType.getElementTypeBitWidth() * numElements) / 8;
+        (dstMemrefType.getElementTypeBitWidth() * dstElements) / 8;
     // bypass L1 is only supported for byte sizes of 16, we drop the hint
     // otherwise.
     UnitAttr bypassL1 =
         sizeInBytes == 16 ? adaptor.getBypassL1Attr() : UnitAttr();
-    rewriter.create<NVVM::CpAsyncOp>(
-        loc, dstPtr, scrPtr, rewriter.getI32IntegerAttr(sizeInBytes), bypassL1);
+
+    // When the optional SrcElements argument is present, the source (global
+    // memory) of CpAsyncOp is read only for SrcElements number of elements. The
+    // rest of the DstElements in the destination (shared memory) are filled
+    // with zeros.
+    if (op.getSrcElements())
+      emitCpAsyncOpZfillAsm(loc, dstPtr, scrPtr,
+                            rewriter.create<LLVM::ConstantOp>(
+                                loc, rewriter.getI32Type(),
+                                rewriter.getI32IntegerAttr(sizeInBytes)),
+                            adaptor.getSrcElements(), srcMemrefType, rewriter);
+
+    // When the optional SrcElements argument is *not* present, the regular
+    // CpAsyncOp is generated. CopyAsyncOp reads bytes from source (global
+    // memory) to fill DstElements number of elements in the destination (shared
+    // memory).
+    else
+      rewriter.create<NVVM::CpAsyncOp>(loc, dstPtr, scrPtr,
+                                       rewriter.getI32IntegerAttr(sizeInBytes),
+                                       bypassL1);
 
     // Drop the result token.
     Value zero = rewriter.create<LLVM::ConstantOp>(
         op->getLoc(), IntegerType::get(op.getContext(), 32),
         rewriter.getI32IntegerAttr(0));
     rewriter.replaceOp(op, zero);
     return success();
   }

mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir

func.func @async_cp_i4(
   // CHECK-DAG: %[[ADDRESSSRC:.*]] = llvm.getelementptr %[[BASESRC]][%[[FI3]]] : (!llvm.ptr<i4>, i64) -> !llvm.ptr<i4>
   // CHECK-DAG: %[[CAST1:.*]] = llvm.bitcast %[[ADDRESSSRC]] : !llvm.ptr<i4> to !llvm.ptr<i8>
   // CHECK-DAG: %[[CAST2:.*]] = llvm.addrspacecast %[[CAST1]] : !llvm.ptr<i8> to !llvm.ptr<i8, 1>
   // CHECK-DAG: nvvm.cp.async.shared.global %[[CAST0]], %[[CAST2]], 16
   %0 = nvgpu.device_async_copy %src[%i, %i], %dst[%i, %i], 32 : memref<128x64xi4> to memref<128x128xi4, 3>
   return %0 : !nvgpu.device.async.token
 }
 
+// -----
+
+// CHECK-LABEL: @async_cp_zfill(
+// CHECK-SAME: %[[IDX:[a-zA-Z0-9_]+]]: index, %[[SRCELEMENTS:[a-zA-Z0-9_]+]]: index)
+func.func @async_cp_zfill(
+  %src: memref<128x128xf32>, %dst: memref<3x16x128xf32, 3>, %i : index, %srcElements : index) {
+
+  // CHECK-DAG: lvm.inline_asm has_side_effects asm_dialect = att "cp.async.cg.shared.global [$0], [$1], $2, $3;\0A", "r,l,n,r" %[[DSTPTR:.*]], %[[SRCPTR:.*]], %[[DSTBYTES:.*]], %[[SRCBYTES:.*]] : (!llvm.ptr<i8, 3>, !llvm.ptr<i8, 1>, i32, i32) -> !llvm.void
+  %0 = nvgpu.device_async_copy %src[%i, %i], %dst[%i, %i, %i], 4, %srcElements {bypassL1}: memref<128x128xf32> to memref<3x16x128xf32, 3>
+  // CHECK: nvvm.cp.async.commit.group
+  %1 = nvgpu.device_async_create_group %0
+  // CHECK: nvvm.cp.async.wait.group 1
+  nvgpu.device_async_wait %1 { numGroups = 1 : i32 }
+
+  return
+}