[mlir][nvvm] Add attribute to nvvm.cpAsyncOp to control l1 bypass

[ThomasRaoux]

Summary

Add attribute to be able to generate the intrinsic version of async copy
generating a copy with l1 bypass. This correspond to
cp.async.cg.shared.global in ptx.

[christopherbate]

LGTM, thanks for adding this.

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

def NVVM_VoteBallotOp :
   }];
   let hasCustomAssemblyFormat = 1;
 }
 
 
 def NVVM_CpAsyncOp : NVVM_Op<"cp.async.shared.global">,
   Arguments<(ins LLVM_i8Ptr_shared:$dst,
                  LLVM_i8Ptr_global:$src,
-                 I32Attr:$size)> {
+                 I32Attr:$size,
+                 OptionalAttr<UnitAttr>:$bypass_l1)> {
   string llvmBuilder = [{
       llvm::Intrinsic::ID id;
       switch ($size) {
         case 4:
           id = llvm::Intrinsic::nvvm_cp_async_ca_shared_global_4;
           break;
         case 8:
           id = llvm::Intrinsic::nvvm_cp_async_ca_shared_global_8;
           break;
         case 16:
+          if(static_cast<bool>($bypass_l1))
+            id = llvm::Intrinsic::nvvm_cp_async_cg_shared_global_16;
+          else
             id = llvm::Intrinsic::nvvm_cp_async_ca_shared_global_16;
           break;
         default:
           llvm_unreachable("unsupported async copy size");
       }
       createIntrinsicCall(builder, id, {$dst, $src});
   }];
   let assemblyFormat = "$dst `,` $src `,` $size attr-dict";
   let hasVerifier = 1;

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

matchAndRewrite(gpu::DeviceAsyncCopyOp op, OpAdaptor adaptor,
     auto srcPointerGlobalType =
         LLVM::LLVMPointerType::get(i8Ty, NVVMMemorySpace::kGlobalMemorySpace);
     scrPtr = rewriter.create<LLVM::AddrSpaceCastOp>(loc, srcPointerGlobalType,
                                                     scrPtr);
     int64_t numElements = adaptor.numElements().getZExtValue();
     int64_t sizeInBytes =
         (dstMemrefType.getElementTypeBitWidth() / 8) * numElements;
     rewriter.create<NVVM::CpAsyncOp>(loc, dstPtr, scrPtr,
-                                     rewriter.getI32IntegerAttr(sizeInBytes));
+                                     rewriter.getI32IntegerAttr(sizeInBytes),
+                                     /*bypassL1=*/UnitAttr());
 
     // 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/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp

return failure(parser.parseOperandList(ops) ||
                                         parser.getNameLoc(), result.operands));
 }
 
 void VoteBallotOp::print(OpAsmPrinter &p) { printNVVMIntrinsicOp(p, *this); }
 
 LogicalResult CpAsyncOp::verify() {
   if (size() != 4 && size() != 8 && size() != 16)
     return emitError("expected byte size to be either 4, 8 or 16.");
+  if (bypass_l1() && size() != 16)
+    return emitError("bypass l1 is only support for 16 bytes copy.");
   return success();
 }
 
 // Given the element type of an operand and whether or not it is an accumulator,
 // this function returns the PTX type (`NVVM::MMATypes`) that corresponds to the
 // operand's element type.
 Optional<mlir::NVVM::MMATypes> MmaOp::inferOperandMMAType(Type operandElType,
                                                           bool isAccumulator) {

mlir/test/Dialect/LLVMIR/invalid.mlir

 func.func @cp_async(%arg0: !llvm.ptr<i8, 3>, %arg1: !llvm.ptr<i8, 1>) {
   // expected-error @below {{expected byte size to be either 4, 8 or 16.}}
   nvvm.cp.async.shared.global %arg0, %arg1, 32
   return
 }
 
 // -----
 
+func.func @cp_async(%arg0: !llvm.ptr<i8, 3>, %arg1: !llvm.ptr<i8, 1>) {
+  // expected-error @below {{bypass l1 is only support for 16 bytes copy.}}
+  nvvm.cp.async.shared.global %arg0, %arg1, 8 {bypass_l1}
+  return
+}
+
+// -----
+
 func.func @gep_struct_variable(%arg0: !llvm.ptr<struct<(i32)>>, %arg1: i32, %arg2: i32) {
   // expected-error @below {{op expected index 1 indexing a struct to be constant}}
   llvm.getelementptr %arg0[%arg1, %arg1] : (!llvm.ptr<struct<(i32)>>, i32, i32) -> !llvm.ptr<i32>
   return
 }
 
 // -----
 

mlir/test/Dialect/LLVMIR/nvvm.mlir

%r = nvvm.wmma.mma %0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15
     : (i32, i32, i32, i32, i32, i32, i32, i32, f32, f32, f32, f32, f32, f32, f32, f32)
     -> !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32)>
   llvm.return %r : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32)>
 }
 
 llvm.func @cp_async(%arg0: !llvm.ptr<i8, 3>, %arg1: !llvm.ptr<i8, 1>) {
 // CHECK:  nvvm.cp.async.shared.global %{{.*}}, %{{.*}}, 16
   nvvm.cp.async.shared.global %arg0, %arg1, 16
+// CHECK:  nvvm.cp.async.shared.global %{{.*}}, %{{.*}}, 16 {bypass_l1}
+  nvvm.cp.async.shared.global %arg0, %arg1, 16 {bypass_l1}
 // CHECK: nvvm.cp.async.commit.group
   nvvm.cp.async.commit.group
 // CHECK: nvvm.cp.async.wait.group 0
   nvvm.cp.async.wait.group 0
   llvm.return
 }
 
 // CHECK-LABEL: llvm.func @ld_matrix

mlir/test/Target/LLVMIR/nvvmir.mlir

 
 llvm.func @cp_async(%arg0: !llvm.ptr<i8, 3>, %arg1: !llvm.ptr<i8, 1>) {
 // CHECK: call void @llvm.nvvm.cp.async.ca.shared.global.4(i8 addrspace(3)* %{{.*}}, i8 addrspace(1)* %{{.*}})
   nvvm.cp.async.shared.global %arg0, %arg1, 4
 // CHECK: call void @llvm.nvvm.cp.async.ca.shared.global.8(i8 addrspace(3)* %{{.*}}, i8 addrspace(1)* %{{.*}})
   nvvm.cp.async.shared.global %arg0, %arg1, 8
 // CHECK: call void @llvm.nvvm.cp.async.ca.shared.global.16(i8 addrspace(3)* %{{.*}}, i8 addrspace(1)* %{{.*}})
   nvvm.cp.async.shared.global %arg0, %arg1, 16
+// CHECK: call void @llvm.nvvm.cp.async.cg.shared.global.16(i8 addrspace(3)* %{{.*}}, i8 addrspace(1)* %{{.*}})
+  nvvm.cp.async.shared.global %arg0, %arg1, 16 {bypass_l1}
 // CHECK: call void @llvm.nvvm.cp.async.commit.group()
   nvvm.cp.async.commit.group
 // CHECK: call void @llvm.nvvm.cp.async.wait.group(i32 0)
   nvvm.cp.async.wait.group 0
   llvm.return
 }
 
 // CHECK-LABEL: @ld_matrix(