[NVGPUToNVVM] Prevent emitting invalid ptx opcode in async.cp lowering

[nicolasvasilache]

Summary

When emitting inline assembly make sure we don't emit invalid ptx opcode

[bondhugula]

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

371

Please add a comment.

391

bypassL1 ? ... : ... should do.

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

303

Should this be lvm or llvm?

[nicolasvasilache]

Similar functionality landed in https://reviews.llvm.org/D148604

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

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,
+                                  UnitAttr bypassL1,
                                   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 *asmStrByPassL1 =
+      "cp.async.cg.shared.global [$0], [$1], $2, $3;\n";
+  const char *asmStr = "cp.async.ca.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,
+      /*asm_string=*/bypassL1 == UnitAttr() ? asmStr : asmStrByPassL1,
       /*constraints=*/asmConstraints, /*has_side_effects=*/true,
       /*is_align_stack=*/false, /*asm_dialect=*/asmDialectAttr,
       /*operand_attrs=*/ArrayAttr());
 }
 
 /// Returns the constraints for the sparse MMA inline assembly instruction.
 static std::string buildMmaSparseAsmConstraintString(unsigned matASize,
                                                      unsigned matBSize,
matchAndRewrite(nvgpu::DeviceAsyncCopyOp op, OpAdaptor adaptor,
     // 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);
+                            adaptor.getSrcElements(), srcMemrefType, bypassL1,
+                            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),

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

 }
 
 // -----
 
 // 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<3>, !llvm.ptr<1>, i32, i32) -> !llvm.void
+  // CHECK: %[[C16:.+]] = llvm.mlir.constant(16 : i32) : i32
+  // CHECK: 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:.*]], %[[C16]], %[[SRCBYTES:.*]] : (!llvm.ptr<3>, !llvm.ptr<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: %[[C4:.+]] = llvm.mlir.constant(4 : i32) : i32
+  // CHECK: lvm.inline_asm has_side_effects asm_dialect = att "cp.async.ca.shared.global [$0], [$1], $2, $3;\0A", "r,l,n,r" %[[DSTPTR:.*]], %[[SRCPTR:.*]], %[[C4]], %[[SRCBYTES:.*]] : (!llvm.ptr<3>, !llvm.ptr<1>, i32, i32) -> !llvm.void
+  %1 = nvgpu.device_async_copy %src[%i, %i], %dst[%i, %i, %i], 1, %srcElements {bypassL1}: memref<128x128xf32> to memref<3x16x128xf32, 3>
   // CHECK: nvvm.cp.async.commit.group
-  %1 = nvgpu.device_async_create_group %0
+  %2 = nvgpu.device_async_create_group %0, %1
   // CHECK: nvvm.cp.async.wait.group 1
-  nvgpu.device_async_wait %1 { numGroups = 1 : i32 }
+  nvgpu.device_async_wait %2 { numGroups = 1 : i32 }
 
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @mma_sp_sync_f16_16832(
 func.func @mma_sp_sync_f16_16832(%arg0: vector<4x2xf16>,