[mlir][nvvm] Add NVVMToLLVM Pass

[guraypp]

Summary

It introduces an NVVMToLLVM Pass and a BasicPtxBuilderOpInterface interface. The Pass performs pattern matching on all the NVVM Ops that implement the BasicPtxBuilderOpInterface interface to generate LLVM Inline Assembly Ops.

The BasicPtxBuilderOpInterface interface is utilized in the convert-nvvm-to-llvm pass, which lowers Ops that support this interface to inline assembly Ops. The interface provides several methods that are used for this lowering.

The getPtx method returns PTX code. The hasSideEffect method is used to determine whether the op has any side effects on the memory. The hasIntrinsic method indicates whether the operation has intrinsic support in LLVM. This is particularly useful for Ops that don't have intrinsic support for each case. The getAsmValues method returns the arguments to be passed to the PTX code. The order of arguments starts with the results and they are used for write operations, followed by the operands and attributes.

Example:

If we have the following Op definition that returns PTX code through getPtx:

tablegen
def NVVM_MBarrierArriveExpectTxOp : NVVM_Op<\"mbarrier.arrive.expect_tx\",
                    [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>]>,
  Results<(outs LLVM_Type:$res)>, Arguments<(ins LLVM_i64ptr_any:$addr, I32:$txcount)> {
  ...
  let extraClassDefinition = [{
    const char* $cppClass::getPtx() { return \"mbarrier.arrive.expect_tx.b64 %0, [%1], %2;\"; }
  }\];
}

The NVVM Op will look like below:

mlir
  %0 = nvvm.mbarrier.arrive.expect_tx %barrier, %txcount : !llvm.ptr, i32 -> i32

The convert-nvvm-to-llvm Pass generates the following PTX code, while keeping the order of arguments the same. The read/write modifiers are set based on the input and result types.

mlir
  %0 = llvm.inline_asm has_side_effects asm_dialect = att "mbarrier.arrive.expect_tx.b64 %0, [%1], %2;", "=r,l,r" %arg0, %arg1 : (!llvm.ptr, i32) -> i32

[guraypp]

add constant support to inline assembly generator

[guraypp]

rebase

[guraypp]

fix typo

[guraypp]

Introduced BasicPtxBuilderOpInterface that automatically builds PTX based on Op.

[guraypp]

Improve the interface

[guraypp]

Implement it with PtxBuilder Interface in
https://reviews.llvm.org/D154060

[guraypp]

rebase

[ThomasRaoux]

nvvm is supposed to match 1 to 1 to nvvm llvm intrinsics so having them lower to inline assembly doesn't make sense to me. Either an intrinsic exist and the nvvm op should lower to it or there wouldn't be an nvvm op.
It seems like this would diverge from the original design of nvvm dialect? Were there some discussions about this?

[guraypp]

In D154060#4474069, @ThomasRaoux wrote:

nvvm is supposed to match 1 to 1 to nvvm llvm intrinsics so having them lower to inline assembly doesn't make sense to me. Either an intrinsic exist and the nvvm op should lower to it or there wouldn't be an nvvm op.
It seems like this would diverge from the original design of nvvm dialect? Were there some discussions about this?

Thanks for the review! I am definitely open to hearing your recommendations and suggestions here.

The new nvvm Ops are designed to perfectly match the potential llvm intrinsic, maintaining a 1:1 correspondence. Both this pass and LLVM intrinsics operate at the same level. It's worth noting that even in LLVM, intrinsics are implemented using inline assembly, which aligns with my thought process.

In my understanding, nvvm serves as the platform for generating device assembly or intrinsic, while the nvgpu dialect acts as a bridge from vector to nvvm. So nvgpu is high level dialect to generate a single line of PTX, and nvvm seems the right place to me. So I put this pass in there. This is of course my thinking, your input here is very valuable.

How do you think we can proceed here? What's your recommendations?

[ThomasRaoux]

In D154060#4474916, @guraypp wrote:

In D154060#4474069, @ThomasRaoux wrote:

nvvm is supposed to match 1 to 1 to nvvm llvm intrinsics so having them lower to inline assembly doesn't make sense to me. Either an intrinsic exist and the nvvm op should lower to it or there wouldn't be an nvvm op.
It seems like this would diverge from the original design of nvvm dialect? Were there some discussions about this?

Thanks for the review! I am definitely open to hearing your recommendations and suggestions here.

The new nvvm Ops are designed to perfectly match the potential llvm intrinsic, maintaining a 1:1 correspondence. Both this pass and LLVM intrinsics operate at the same level.

The part that looks odd to me is that until those intrinsics exist it is hard to be sure that what we are adding to nvvm dialect matches well.

It's worth noting that even in LLVM, intrinsics are implemented using inline assembly, which aligns with my thought process.

Do you have examples? I haven't see it in the codebase so far.

In my understanding, nvvm serves as the platform for generating device assembly or intrinsic, while the nvgpu dialect acts as a bridge from vector to nvvm. So nvgpu is high level dialect to generate a single line of PTX, and nvvm seems the right place to me. So I put this pass in there. This is of course my thinking, your input here is very valuable.

nvvm as llvm or rocdl dialect is meant to match as closely as possible llvm IR in it's design. Moving away from that sounds like a non negligible change to me. What I'm wondering is whether we even need this intermediate step.

How do you think we can proceed here? What's your recommendations?

My recommendation would be to directly generate inline assembly when converting from nvgpu/gpu to nvvm as it is done for existing cases using inline assembly.

[guraypp]

The part that looks odd to me is that until those intrinsics exist it is hard to be sure that what we are adding to nvvm dialect matches well.

I thought as long as nvvm op matches with PTX, it also matches with the potential intrinsic.

Do you have examples? I haven't see it in the codebase so far.

See an example below. I mean calling an llvm intrinsic vs blob of inline assembly are same for the llvm or its optimizer. Please correct me if I am wrong.
https://github.com/llvm/llvm-project/blob/5e807c38bf9b5c1528012ec9003953b352f99c79/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td#L76-L81

My recommendation would be to directly generate inline assembly when converting from nvgpu/gpu to nvvm as it is done for existing cases using inline assembly.

In this case, nvgpu will be as low as PTX level. Would it be desirable?

Thanks for the review again.

Maybe I should create a discussion in discourse :)

[ThomasRaoux]

In D154060#4475381, @guraypp wrote:

The part that looks odd to me is that until those intrinsics exist it is hard to be sure that what we are adding to nvvm dialect matches well.

I thought as long as nvvm op matches with PTX, it also matches with the potential intrinsic.

Do you have examples? I haven't see it in the codebase so far.

See an example below. I mean calling an llvm intrinsic vs blob of inline assembly are same for the llvm or its optimizer. Please correct me if I am wrong.
https://github.com/llvm/llvm-project/blob/5e807c38bf9b5c1528012ec9003953b352f99c79/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td#L76-L81

This is the assembler code for MIR instructions, I don't think it is really comparable. This is basically code that convert MIR to the equivalent assembly, this is much lower than intrinsics and obviously the only way to emit PTX assembly. I don't think you can compare that to inline assembly.

My recommendation would be to directly generate inline assembly when converting from nvgpu/gpu to nvvm as it is done for existing cases using inline assembly.

In this case, nvgpu will be as low as PTX level. Would it be desirable?

That shouldn't change the level of nvgpu and shouldn't affect the design or level of abstraction of those ops. It just means when lowering from nvgpu to nvvm instead of emitting nvvm op that gets then converted to inline assembly you directly emit inline assembly. That also doesn't mean you cannot build abstraction in your lowering code.

Thanks for the review again.

Maybe I should create a discussion in discourse :)

It's up to you, in my opinion if we decide to diverge from the original idea that nvvm is a 1 to 1 mapping to intrinsics this is probably something that should be discussed as this is not the original intent of the design. I'm not saying it is a good or a bad idea but we should have a clear direction on what the dialect should and shouldn't represent.

[ThomasRaoux]

To summarize my main point is that having an intermediate step before LLVM inline assembly op seems unnecessary and as far as I know is not how nvvm was meant to be used. If you think there is a need for this it would be interesting to mention it, otherwise my guess is that it is actually simpler to skip the intermediate representation.

[guraypp]

In D154060#4475942, @ThomasRaoux wrote:

To summarize my main point is that having an intermediate step before LLVM inline assembly op seems unnecessary and as far as I know is not how nvvm was meant to be used. If you think there is a need for this it would be interesting to mention it, otherwise my guess is that it is actually simpler to skip the intermediate representation.

I think nvvm is not unnecessary. Allow me to make a final attempt to convince you. Having BasicPtxBuilderOpInterface in nvvm automates inline assembly generation without writing C++ while nvgpu remains important as it serves as a bridge from memref->llvm lowering.

It is particularly beneficial for ops like cp.async with half nvptx support. See the example below, I can generate intrinsic and PTX without touching C++ (D154345). Previosuly, cp.async calls emitCpAsyncOpZfillAsm that's boilerplate code. Don't you think BasicPtxBuilderOpInterface simplifies a lot?

I understand the design concern of nvvm. If you think this work is beneficial, I start a discussion to relax the design (or I can split nvvm.td file into nvvm.td and nvvm_inlineasm.td). But if you think there is no value, I can work on retargeting the interface on nvgpu dialect.

def NVVM_CpAsyncOp : NVVM_Op<"cp.async.shared.global", [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>]>,
  ...
  string llvmBuilder = [{
      ...
      createIntrinsicCall(builder, id, {$dst, $src});
  }];
  let extraClassDeclaration = [{
    bool canBuildPtx() { if(getCpSize()) return true; return false; }

    void getAsmValues(RewriterBase &rewriter, 
        llvm::SmallVectorImpl<std::pair<mlir::Value, mlir::NVVM::PTXRegisterMod>> &asmValues) {
      asmValues.push_back({getDst(), PTXRegisterMod::Read});
      asmValues.push_back({getSrc(), PTXRegisterMod::Read});
      asmValues.push_back({makeConstant(rewriter, getSize()), PTXRegisterMod::Read});
      asmValues.push_back({getCpSize(), PTXRegisterMod::Read});
    }        
  }];
  let extraClassDefinition = [{        
    const char* $cppClass::getPtx() { 
      if(getModifier() == NVVM::LoadCacheModifierKind::CG)
        return "cp.async.cg.shared.global [%0], [%1], %2, %3;\n"; 
      if(getModifier() == NVVM::LoadCacheModifierKind::CA)
        return "cp.async.ca.shared.global [%0], [%1], %2, %3;\n";        
      llvm_unreachable("unsupported cache modifier");      
    }
  }];
}

If op has no intrinsic, generating PTX is even simpler and defining getPtx is enough.

def NVVM_MBarrierTryWaitParitySharedOp : NVVM_Op<"mbarrier.try_wait.parity.shared", 
                    [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>]>,
  Results<(outs LLVM_Type:$res)>,
  Arguments<(ins LLVM_i64ptr_shared:$addr, LLVM_Type:$token)> {  
  let assemblyFormat = "$addr `,` $token attr-dict `:` type(operands) `->` type($res)";
  let extraClassDefinition = [{
    const char* $cppClass::getPtx() {
      return "{\n\t"
              ".reg .pred P1; \n\t"
              "mbarrier.try_wait.parity.shared.b64 P1, [%1], %2; \n\t"
              "selp.b32 %0, 1, 0, P1; \n\t"
              "}"; 
    }
  }];
}

[guraypp]

I put a draft code here D154624. I put the BasicPtxBuilder as a class and used nvgpu. In this case, building ptx happens in C++. I am not sure it is any better.

[ThomasRaoux]

In D154060#4477294, @guraypp wrote:

In D154060#4475942, @ThomasRaoux wrote:

To summarize my main point is that having an intermediate step before LLVM inline assembly op seems unnecessary and as far as I know is not how nvvm was meant to be used. If you think there is a need for this it would be interesting to mention it, otherwise my guess is that it is actually simpler to skip the intermediate representation.

I think nvvm is not unnecessary. Allow me to make a final attempt to convince you. Having BasicPtxBuilderOpInterface in nvvm automates inline assembly generation without writing C++ while nvgpu remains important as it serves as a bridge from memref->llvm lowering.

It is particularly beneficial for ops like cp.async with half nvptx support. See the example below, I can generate intrinsic and PTX without touching C++ (D154345). Previosuly, cp.async calls emitCpAsyncOpZfillAsm that's boilerplate code. Don't you think BasicPtxBuilderOpInterface simplifies a lot?

I understand the design concern of nvvm. If you think this work is beneficial, I start a discussion to relax the design (or I can split nvvm.td file into nvvm.td and nvvm_inlineasm.td). But if you think there is no value, I can work on retargeting the interface on nvgpu dialect.

def NVVM_CpAsyncOp : NVVM_Op<"cp.async.shared.global", [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>]>,
  ...
  string llvmBuilder = [{
      ...
      createIntrinsicCall(builder, id, {$dst, $src});
  }];
  let extraClassDeclaration = [{
    bool canBuildPtx() { if(getCpSize()) return true; return false; }

    void getAsmValues(RewriterBase &rewriter, 
        llvm::SmallVectorImpl<std::pair<mlir::Value, mlir::NVVM::PTXRegisterMod>> &asmValues) {
      asmValues.push_back({getDst(), PTXRegisterMod::Read});
      asmValues.push_back({getSrc(), PTXRegisterMod::Read});
      asmValues.push_back({makeConstant(rewriter, getSize()), PTXRegisterMod::Read});
      asmValues.push_back({getCpSize(), PTXRegisterMod::Read});
    }        
  }];
  let extraClassDefinition = [{        
    const char* $cppClass::getPtx() { 
      if(getModifier() == NVVM::LoadCacheModifierKind::CG)
        return "cp.async.cg.shared.global [%0], [%1], %2, %3;\n"; 
      if(getModifier() == NVVM::LoadCacheModifierKind::CA)
        return "cp.async.ca.shared.global [%0], [%1], %2, %3;\n";        
      llvm_unreachable("unsupported cache modifier");      
    }
  }];
}

If op has no intrinsic, generating PTX is even simpler and defining getPtx is enough.

def NVVM_MBarrierTryWaitParitySharedOp : NVVM_Op<"mbarrier.try_wait.parity.shared", 
                    [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>]>,
  Results<(outs LLVM_Type:$res)>,
  Arguments<(ins LLVM_i64ptr_shared:$addr, LLVM_Type:$token)> {  
  let assemblyFormat = "$addr `,` $token attr-dict `:` type(operands) `->` type($res)";
  let extraClassDefinition = [{
    const char* $cppClass::getPtx() {
      return "{\n\t"
              ".reg .pred P1; \n\t"
              "mbarrier.try_wait.parity.shared.b64 P1, [%1], %2; \n\t"
              "selp.b32 %0, 1, 0, P1; \n\t"
              "}"; 
    }
  }];
}

I understand the motivation and I agree it has some advantages and it allows having an extra level where both intrinsics and inline assembly share a representation. I don't know if there is a use for this extra level at the time which is why I was saying it is probably not necessary. My intuition would be that it is less code and less maintenance to skip the extra level of representation for inline assembly. I'm not sure I follow how this makes the inline ptx emission simpler. Whether the code is in tablegen or not it is still c++ code that needs to emit either the intrinsic or inline PTX? I would assume that the code would be the same emitting directly in conversion to nvvm.

If you think this is better I'm not going to block your work, I just wanted to point out that this doesn't match my mental model of nvvm. I'm not extremely active in this code recently so it probably make sense for someone else to comment more in the area if needed (Alex has a better background than me in LLVM Dialect). My advice in term of simplicity would be to skip the extra op representation but feel free to ignore it. The PTX builder looks neat.

[ThomasRaoux]

In D154060#4479680, @ThomasRaoux wrote:

In D154060#4477294, @guraypp wrote:

In D154060#4475942, @ThomasRaoux wrote:

It is particularly beneficial for ops like cp.async with half nvptx support. See the example below, I can generate intrinsic and PTX without touching C++ (D154345). Previosuly, cp.async calls emitCpAsyncOpZfillAsm that's boilerplate code. Don't you think BasicPtxBuilderOpInterface simplifies a lot?

I see your point about levering tablgen to simplify ptx declaration. I assume something similar could be done in C++ but this is definitely a cool representation. I agree this is an advantage of this approach.

[nicolasvasilache]

I have not followed the full discussion here but consider the following case:

1. some op exists at the nvgpu level on higher-level MLIR types
2. it lowers to a bunch of ops in the NVVM dialect
3. there is not yet an intrinsic representation in LLVM, there should be one, but we have no ETA on this

Isn't the NVVM dialect the right place to put the op and a lowering to PTX ?

Otherwise do we need to start distinguishing in NVGPU ops that should become LLVM intrinsics in the fullness of time and quarantine them for a potential future move to NVVM once the LLVM intrinsic is available ?

Maybe I missed the point of the discussion, please lmk if I need to dig deeper to better understand what the fundamental issues are here.

[ThomasRaoux]

In D154060#4481135, @nicolasvasilache wrote:

I have not followed the full discussion here but consider the following case:

1. some op exists at the nvgpu level on higher-level MLIR types
2. it lowers to a bunch of ops in the NVVM dialect
3. there is not yet an intrinsic representation in LLVM, there should be one, but we have no ETA on this

Isn't the NVVM dialect the right place to put the op and a lowering to PTX ?

Or you can directly LLVM inline assembly ops where you would emit the NVVM that doesn't exist because there is no intrinsic (as it is currently done for cases that don't have intrinsics). It sill happens in NVGPU to NVVM but you don't create an intermediate NVVM op that then needs to be lowered to LLVM dialect.

Otherwise do we need to start distinguishing in NVGPU ops that should become LLVM intrinsics in the fullness of time and quarantine them for a potential future move to NVVM once the LLVM intrinsic is available ?

We definitely don't want to do that. My comments were unrelated to NVGPU dialect and shouldn't affects what ops are added there.

Maybe I missed the point of the discussion, please lmk if I need to dig deeper to better understand what the fundamental issues are here.

I don't think we need much more discussions on that. I just meant to give my opinion and it is totally fine to continue with what you have.

[nicolasvasilache]

In D154060#4481172, @ThomasRaoux wrote:

In D154060#4481135, @nicolasvasilache wrote:

I have not followed the full discussion here but consider the following case:

1. some op exists at the nvgpu level on higher-level MLIR types
2. it lowers to a bunch of ops in the NVVM dialect
3. there is not yet an intrinsic representation in LLVM, there should be one, but we have no ETA on this

Isn't the NVVM dialect the right place to put the op and a lowering to PTX ?

Or you can directly LLVM inline assembly ops where you would emit the NVVM that doesn't exist because there is no intrinsic (as it is currently done for cases that don't have intrinsics). It sill happens in NVGPU to NVVM but you don't create an intermediate NVVM op that then needs to be lowered to LLVM dialect.

Ah yes ok, that is where NGPU kicks in, I had glanced over the discussion and mistakenly thought we talking about the NVGPU dialect.

Otherwise do we need to start distinguishing in NVGPU ops that should become LLVM intrinsics in the fullness of time and quarantine them for a potential future move to NVVM once the LLVM intrinsic is available ?

We definitely don't want to do that. My comments were unrelated to NVGPU dialect and shouldn't affects what ops are added there.

+1, sorry for the confusion!

Maybe I missed the point of the discussion, please lmk if I need to dig deeper to better understand what the fundamental issues are here.

I don't think we need much more discussions on that. I just meant to give my opinion and it is totally fine to continue with what you have.

I'd certainly want the better PTX builder and nice inference of the magic r m etc.

I have no strong opinion on the NVVM ops themselves.
It seems a bit closer to the final landing place in NVVM and then the switch from InlineAsm to LLVM intrinsic can happen locally within 1 file when ready but it is not a big blocker to also update the Conversion at the same time.

[nicolasvasilache]

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

164

makeConstantI32 please ? since this is ahrdcoded ?

174

Can you expand the description here to say that the default implementation emits results, inputs and then attributes?

[nicolasvasilache]

Upon deeper inspection, I se great value in the way this is built, we can iterate to further improvements if needed in the future, let's get this rolling for now!

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

140

this is off it seems

149

Return whether the operation has memory side effects.

408

it is really great that we can actually just see the ptx with a "".format-like behavior.
This greatly simplifies the understanding vs building C++ by hand and reduces surprise.

Nice work!

mlir/lib/Conversion/NVVMToLLVM/NVVMToLLVM.cpp

72

Add TODO: f8,f16,bf16 etc ?

188

nit: nl

mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir

25

nice!

[guraypp]

address comments

[guraypp]

fix debug printer (errs -> dbgs)

mlir/include/mlir/Conversion/NVVMToLLVM/NVVMToLLVM.h

+//===- NVVMTOLLVMPass.h - Convert NVVM to LLVM dialect --------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_CONVERSION_NVVMTOLLVM_NVVMTOLLVMPASS_H_
+#define MLIR_CONVERSION_NVVMTOLLVM_NVVMTOLLVMPASS_H_
+
+#include <memory>
+
+namespace mlir {
+
+class LLVMTypeConverter;
+class RewritePatternSet;
+class Pass;
+
+#define GEN_PASS_DECL_CONVERTNVVMTOLLVMPASS
+#include "mlir/Conversion/Passes.h.inc"
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_NVVMTOLLVM_NVVMTOLLVMPASS_H_

mlir/include/mlir/Conversion/Passes.h

 #include "mlir/Conversion/LinalgToStandard/LinalgToStandard.h"
 #include "mlir/Conversion/MathToFuncs/MathToFuncs.h"
 #include "mlir/Conversion/MathToLLVM/MathToLLVM.h"
 #include "mlir/Conversion/MathToLibm/MathToLibm.h"
 #include "mlir/Conversion/MathToSPIRV/MathToSPIRVPass.h"
 #include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
 #include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRVPass.h"
 #include "mlir/Conversion/NVGPUToNVVM/NVGPUToNVVM.h"
+#include "mlir/Conversion/NVVMToLLVM/NVVMToLLVM.h"
 #include "mlir/Conversion/OpenACCToSCF/ConvertOpenACCToSCF.h"
 #include "mlir/Conversion/OpenMPToLLVM/ConvertOpenMPToLLVM.h"
 #include "mlir/Conversion/PDLToPDLInterp/PDLToPDLInterp.h"
 #include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"
 #include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
 #include "mlir/Conversion/SCFToGPU/SCFToGPUPass.h"
 #include "mlir/Conversion/SCFToOpenMP/SCFToOpenMP.h"
 #include "mlir/Conversion/SCFToSPIRV/SCFToSPIRVPass.h"

mlir/include/mlir/Conversion/Passes.td

def ConvertMemRefToSPIRV : Pass<"convert-memref-to-spirv"> {
   let options = [
     Option<"boolNumBits", "bool-num-bits",
            "int", /*default=*/"8",
            "The number of bits to store a boolean value">
   ];
 }
 
 //===----------------------------------------------------------------------===//
+// NVVMToLLVM
+//===----------------------------------------------------------------------===//
+
+def ConvertNVVMToLLVMPass : Pass<"convert-nvvm-to-llvm"> {
+  let summary = "Convert NVVM dialect to LLVM dialect";
+  let description = [{
+    This pass generates inline assembly for the NVVM ops which is not 
+    implemented in LLVM core.
+  }];
+  let dependentDialects = [
+    "NVVM::NVVMDialect",
+  ];  
+}
+
+//===----------------------------------------------------------------------===//
 // NVGPUToNVVM
 //===----------------------------------------------------------------------===//
 
 def ConvertNVGPUToNVVMPass : Pass<"convert-nvgpu-to-nvvm"> {
   let summary = "Convert NVGPU dialect to NVVM dialect";
   let description = [{
     This pass converts supported NVGPU ops to NVVM dialect intrinsics.
   }];

mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt

 add_public_tablegen_target(MLIRLLVMIntrinsicConversionsIncGen)
 
 add_mlir_dialect(NVVMOps nvvm)
 add_mlir_doc(NVVMOps NVVMDialect Dialects/ -gen-dialect-doc -dialect=nvvm)
 set(LLVM_TARGET_DEFINITIONS NVVMOps.td)
 mlir_tablegen(NVVMConversions.inc -gen-llvmir-conversions)
 mlir_tablegen(NVVMOpsEnums.h.inc -gen-enum-decls)
 mlir_tablegen(NVVMOpsEnums.cpp.inc -gen-enum-defs)
+mlir_tablegen(NVVMOpsInterface.h.inc -gen-op-interface-decls)
+mlir_tablegen(NVVMOpsInterface.cpp.inc -gen-op-interface-defs)
 mlir_tablegen(NVVMOpsAttributes.h.inc -gen-attrdef-decls -attrdefs-dialect=nvvm)
 mlir_tablegen(NVVMOpsAttributes.cpp.inc -gen-attrdef-defs -attrdefs-dialect=nvvm)
 add_public_tablegen_target(MLIRNVVMConversionsIncGen)
 
 add_mlir_dialect(ROCDLOps rocdl)
 add_mlir_doc(ROCDLOps ROCDLDialect Dialects/ -gen-dialect-doc -dialect=rocdl)
 set(LLVM_TARGET_DEFINITIONS ROCDLOps.td)
 mlir_tablegen(ROCDLConversions.inc -gen-llvmir-conversions)
 add_public_tablegen_target(MLIRROCDLConversionsIncGen)

mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h

 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "llvm/IR/IntrinsicsNVPTX.h"
 
 #include "mlir/Dialect/LLVMIR/NVVMOpsEnums.h.inc"
 
 namespace mlir {
 namespace NVVM {
 
+#include "mlir/Dialect/LLVMIR/NVVMOpsInterface.h.inc"
+
 /// NVVM memory space identifiers.
 enum NVVMMemorySpace {
   /// Global memory space identifier.
   kGlobalMemorySpace = 1,
   /// Shared memory space identifier.
   kSharedMemorySpace = 3
 };
 

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

 //===----------------------------------------------------------------------===//
 
 class NVVM_Attr<string attrName, string attrMnemonic, list<Trait> traits = []>
     : AttrDef<NVVM_Dialect, attrName, traits> {
   let mnemonic = attrMnemonic;
 }
 
 //===----------------------------------------------------------------------===//
+// Basic PTX Builder Interface
+//===----------------------------------------------------------------------===//
+
+// https://docs.nvidia.com/cuda/inline-ptx-assembly/index.html#parameters
+def Read : I32EnumAttrCase<"Read", 0, "read">;
+def Write : I32EnumAttrCase<"Write", 2, "write">;
+def ReadWrite : I32EnumAttrCase<"ReadWrite", 1, "readwrite">;
+
+def PTXRegisterMod : I32EnumAttr<"PTXRegisterMod", 
+  "Register read/write modifier to build cosntraint string for PTX inline",
+  [Read, Write, ReadWrite]> {
+  let cppNamespace = "::mlir::NVVM";
+}
+
+
+def BasicPtxBuilderOpInterface : OpInterface<"BasicPtxBuilderInterface"> {
+  let description = [{
+    Interface to generate inline assembly with PTX for basic operations.
+
+    It uses `getPtx` interface function to read PTX code. The `hasSideEffect`
+    is used to check if the PTX code has any side effect on the memory.
+
+    The order of PTX arguments is started from the results, followed by the 
+    operands. The order of arguments is the same as the order of the operands.
+
+    Example:
+    
+    If we have following Op definition:
+    ```tablegen
+      def NVVM_MyOp : NVVM_Op<\"myop\", 
+          [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>]>,
+          Results<(outs LLVM_Type:$res)>,
+          Arguments<(ins LLVM_i64ptr_any:$arg1, I32:$arg2)> {          
+          let extraClassDefinition = [{
+           bool $cppClass::hasSideEffect() { return false; }
+           const char* $cppClass::getPtx() { return\"ptx.code %0, [%1], %2;\";}
+          }\]
+    ```
+
+    The NVVM Op will look like below:
+    ```mlir
+      nvvm.op %1, %2 : !llvm.ptr, i32 -> i32
+    ```
+
+    The `convert-nvvm-to-llvm` Pass returns the PTX code below. The order of 
+    arguments are kept the same. The read/write modifiers are set based on the
+     input and result types.
+    ```mlir
+      llvm.inline_asm asm_dialect = att 
+          "ptx.code %0, [%1], %2;", "=r,l,r" %1, %2 : (!llvm.ptr, i32) -> i32
+    ```
+  }];
+  let methods = [
+    InterfaceMethod<
+        /*desc=*/[{
+          Returns a list of iterator types that describe the number of loops.
+        }],
+        /*retType=*/"bool",
+        /*methodName=*/"canBuildPtx",
+        /*args=*/(ins),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/"return true;"
+      >,
+    InterfaceMethod<
+        /*desc=*/[{ Returns memory side effect of the PTX code. }],
+        /*retType=*/"bool",
+        /*methodName=*/"hasSideEffect",
+        /*args=*/(ins),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/"return true;"
+      >,
+    InterfaceMethod<
+        /*desc=*/[{ Returns PTX code. }],
+        /*retType=*/"const char*",
+        /*methodName=*/"getPtx"
+      >,
+    InterfaceMethod<
+        /*desc=*/[{Generate constant value.}],
+        /*retType=*/"::mlir::Value",
+        /*methodName=*/"makeConstant",
+        /*args=*/(ins "::mlir::RewriterBase &":$rewriter, "unsigned" : $val),
+        /*methodBody=*/"",
+        /*defaultImpl=*/ [{
+            mlir::Operation* op = $_op;
+            return rewriter.create<LLVM::ConstantOp>(
+              op->getLoc(), rewriter.getIntegerType(32), val);
+        }]
+     >,
+     InterfaceMethod<
+         /*desc=*/[{ Returns arguments to pass PTX code. }],
+         /*retType=*/"void",
+         /*methodName=*/"getAsmValues",
+         /*args=*/(ins "::mlir::RewriterBase &":$rewriter, 
+         "llvm::SmallVectorImpl<std::pair<mlir::Value, mlir::NVVM::PTXRegisterMod>>&" : $asmValues),
+         /*methodBody=*/"",
+         /*defaultImpl=*/ [{         
+           mlir::Operation* op = $_op;
+           for (auto val : op->getResults()) 
+            asmValues.push_back({val, mlir::NVVM::PTXRegisterMod::Write});
+           for (auto val : op->getOperands()) 
+            asmValues.push_back({val, mlir::NVVM::PTXRegisterMod::Read});
+           for (auto attr : op->getAttrs()) {
+            if (auto intAttr = dyn_cast<mlir::IntegerAttr>(attr.getValue())) {
+             Value val = makeConstant(rewriter, intAttr.getInt());
+             asmValues.push_back({val ,mlir::NVVM::PTXRegisterMod::Read});
+             }
+           }
+         }]
+       >
+  ];
+}
+
+//===----------------------------------------------------------------------===//
 // NVVM intrinsic operations
 //===----------------------------------------------------------------------===//
 
 class NVVM_IntrOp<string mnem, list<Trait> traits,
                   int numResults>
   : LLVM_IntrOpBase<NVVM_Dialect, mnem, "nvvm_" # !subst(".", "_", mnem),
                     /*list<int> overloadedResults=*/[],
                     /*list<int> overloadedOperands=*/[],
def NVVM_MBarrierArriveNocompleteSharedOp : NVVM_Op<"mbarrier.arrive.nocomplete.shared">,
   Results<(outs LLVM_Type:$res)>,
   Arguments<(ins LLVM_i64ptr_shared:$addr, I32:$count)> {
   string llvmBuilder = [{
       $res = createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mbarrier_arrive_noComplete_shared, {$addr, $count});
   }];
   let assemblyFormat = "$addr `,` $count attr-dict `:` type(operands) `->` type($res)";
 }
 
+def NVVM_MBarrierArriveExpectTxOp : NVVM_Op<"mbarrier.arrive.expect_tx",
+                    [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>]>,
+  Results<(outs LLVM_Type:$res)>,
+  Arguments<(ins LLVM_i64ptr_any:$addr, I32:$txcount)> {
+  let assemblyFormat = "$addr `,` $txcount attr-dict `:` type(operands) `->` type($res)";
+  let extraClassDefinition = [{
+    const char* $cppClass::getPtx() { return "mbarrier.arrive.expect_tx.b64 %0, [%1], %2;"; }
+  }];
+}
+
+def NVVM_MBarrierArriveExpectTxSharedOp : NVVM_Op<"mbarrier.arrive.expect_tx.shared", 
+                    [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>]>,
+  Results<(outs LLVM_Type:$res)>,
+  Arguments<(ins LLVM_i64ptr_shared:$addr, I32:$txcount)> {    
+  let assemblyFormat = "$addr `,` $txcount attr-dict `:` type(operands) `->` type($res)";
+  let extraClassDefinition = [{
+    const char* $cppClass::getPtx() { return "mbarrier.arrive.expect_tx.shared.b64 %0, [%1], %2;"; }    
+  }];
+}
+
+def NVVM_MBarrierTryWaitParityOp : NVVM_Op<"mbarrier.try_wait.parity", 
+                    [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>]>,
+  Results<(outs LLVM_Type:$res)>,
+  Arguments<(ins LLVM_i64ptr_any:$addr, LLVM_Type:$token)> {
+  let assemblyFormat = "$addr `,` $token attr-dict `:` type(operands) `->` type($res)";
+  let extraClassDefinition = [{
+    const char* $cppClass::getPtx() {
+      return "{\n\t"
+              ".reg .pred P1; \n\t"
+              "mbarrier.try_wait.parity.b64 P1, [%1], %2; \n\t"
+              "selp.b32 %0, 1, 0, P1; \n\t"
+              "}"; 
+    }
+  }];
+}
+
+def NVVM_MBarrierTryWaitParitySharedOp : NVVM_Op<"mbarrier.try_wait.parity.shared", 
+                    [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>]>,
+  Results<(outs LLVM_Type:$res)>,
+  Arguments<(ins LLVM_i64ptr_shared:$addr, LLVM_Type:$token)> {  
+  let assemblyFormat = "$addr `,` $token attr-dict `:` type(operands) `->` type($res)";
+  let extraClassDefinition = [{
+    const char* $cppClass::getPtx() {
+      return "{\n\t"
+              ".reg .pred P1; \n\t"
+              "mbarrier.try_wait.parity.shared.b64 P1, [%1], %2; \n\t"
+              "selp.b32 %0, 1, 0, P1; \n\t"
+              "}"; 
+    }
+  }];
+}
+
 def NVVM_MBarrierTestWaitOp : NVVM_Op<"mbarrier.test.wait">,
   Results<(outs LLVM_Type:$res)>,
   Arguments<(ins LLVM_i64ptr_any:$addr, LLVM_Type:$token)> {
   string llvmBuilder = [{
       $res = createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mbarrier_test_wait, {$addr, $token});
   }];
   let assemblyFormat = "$addr `,` $token attr-dict `:` type(operands) `->` type($res)";
 }

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(LLVMCommon)
 add_subdirectory(MathToFuncs)
 add_subdirectory(MathToLibm)
 add_subdirectory(MathToLLVM)
 add_subdirectory(MathToSPIRV)
 add_subdirectory(MemRefToLLVM)
 add_subdirectory(MemRefToSPIRV)
 add_subdirectory(NVGPUToNVVM)
+add_subdirectory(NVVMToLLVM)
 add_subdirectory(OpenACCToSCF)
 add_subdirectory(OpenMPToLLVM)
 add_subdirectory(PDLToPDLInterp)
 add_subdirectory(ReconcileUnrealizedCasts)
 add_subdirectory(SCFToControlFlow)
 add_subdirectory(SCFToGPU)
 add_subdirectory(SCFToOpenMP)
 add_subdirectory(SCFToSPIRV)

mlir/lib/Conversion/NVVMToLLVM/CMakeLists.txt

+add_mlir_conversion_library(MLIRNVVMToLLVM
+  NVVMToLLVM.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/NVVMToLLVM
+
+  DEPENDS
+  MLIRConversionPassIncGen
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIRGPUDialect
+  MLIRLLVMCommonConversion
+  MLIRLLVMDialect
+  MLIRNVVMDialect
+  MLIRNVGPUDialect
+  MLIRPass
+  MLIRTransforms
+  )

mlir/lib/Conversion/NVVMToLLVM/NVVMToLLVM.cpp

+//===- NVVMToLLVM.cpp - NVVM to LLVM dialect conversion -----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a translation NVVM ops which is not supported in LLVM
+// core.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/NVVMToLLVM/NVVMToLLVM.h"
+
+#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
+#include "mlir/Conversion/LLVMCommon/Pattern.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
+#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/TypeUtilities.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/raw_ostream.h"
+#include <string>
+
+#define DEBUG_TYPE "nvvm-to-llvm"
+#define DBGS() (llvm::errs() << "[" DEBUG_TYPE "]: ")
+#define DBGSNL() (llvm::errs() << "\n")
+
+namespace mlir {
+#define GEN_PASS_DEF_CONVERTNVVMTOLLVMPASS
+#include "mlir/Conversion/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+using namespace NVVM;
+
+#include "mlir/Dialect/LLVMIR/NVVMOpsInterface.cpp.inc"
+namespace {
+
+class PtxBuilder {
+  Operation *op;
+  PatternRewriter &rewriter;
+  const char *asmStr;
+  SmallVector<Value> asmVals;
+  std::string asmConstraints;
+  bool sideEffects;
+  bool hasResult = false;
+
+  // https://docs.nvidia.com/cuda/inline-ptx-assembly/index.html#constraints
+  char getRegisterType(Value v) {
+    if (v.getDefiningOp<LLVM::ConstantOp>())
+      return 'n';
+    if (v.getType().isInteger(16))
+      return 'h';
+    if (v.getType().isInteger(32))
+      return 'r';
+    if (v.getType().isInteger(64))
+      return 'l';
+    if (v.getType().isF32())
+      return 'f';
+    if (v.getType().isF64())
+      return 'd';
+    if (auto ptr = v.getType().dyn_cast<LLVM::LLVMPointerType>()) {
+      // Shared address spaces is addressed with 32-bit pointers.
+      if (ptr.getAddressSpace() == NVVM::kSharedMemorySpace) {
+        return 'r';
+      }
+      return 'l';
+    }
+    assert(false && "Register type is not handled yet");
+    return ' ';
+  }
+
+public:
+  PtxBuilder(Operation *op, PatternRewriter &rewriter, const char *ptxAsm,
+             bool sideEffects = false)
+      : op(op), rewriter(rewriter), asmStr(ptxAsm), sideEffects(sideEffects) {}
+
+  void insertValue(Value v, PTXRegisterMod itype = PTXRegisterMod::Read) {
+    llvm::raw_string_ostream ss(asmConstraints);
+    if (itype == PTXRegisterMod::Read) {
+      asmVals.push_back(v);
+    } else if (itype == PTXRegisterMod::ReadWrite) {
+      asmVals.push_back(v);
+      ss << "+";
+      hasResult = true;
+    } else if (itype == PTXRegisterMod::Write) {
+      ss << "=";
+      hasResult = true;
+    }
+    ss << getRegisterType(v) << ",";
+    ss.flush();
+  }
+
+  LLVM::InlineAsmOp build() {
+    auto asmDialectAttr =
+        LLVM::AsmDialectAttr::get(op->getContext(), LLVM::AsmDialect::AD_ATT);
+    Type resultType = hasResult ? op->getResult(0).getType()
+                                : LLVM::LLVMVoidType::get(op->getContext());
+
+    // Remove the last comma from the constraints string.
+    if (asmConstraints[asmConstraints.size() - 1] == ',')
+      asmConstraints.pop_back();
+
+    return rewriter.create<LLVM::InlineAsmOp>(
+        op->getLoc(), resultType,
+        /*operands=*/asmVals,
+        /*asm_string=*/asmStr,
+        /*constraints=*/asmConstraints.data(),
+        /*has_side_effects=*/sideEffects,
+        /*is_align_stack=*/false,
+        /*asm_dialect=*/asmDialectAttr,
+        /*operand_attrs=*/ArrayAttr());
+  }
+
+  void buildAndReplaceOp() {
+    LLVM::InlineAsmOp inlineAsmOp = build();
+    LLVM_DEBUG(DBGS() << "\n Generated PTX \n\t" << inlineAsmOp << "\n");
+    if (inlineAsmOp->getNumResults() == op->getNumResults())
+      rewriter.replaceOp(op, inlineAsmOp);
+    else
+      rewriter.eraseOp(op);
+  }
+};
+
+struct PtxLowering
+    : public OpInterfaceRewritePattern<NVVM::BasicPtxBuilderInterface> {
+  using OpInterfaceRewritePattern<
+      NVVM::BasicPtxBuilderInterface>::OpInterfaceRewritePattern;
+
+  PtxLowering(MLIRContext *context, PatternBenefit benefit = 2)
+      : OpInterfaceRewritePattern(context, benefit) {}
+
+  LogicalResult matchAndRewrite(NVVM::BasicPtxBuilderInterface op,
+                                PatternRewriter &rewriter) const override {
+    if (!op.canBuildPtx()) {
+      LLVM_DEBUG(DBGS() << "Ptx Builder does not lower \n\t" << op << "\n");
+      return failure();
+    }
+
+    SmallVector<std::pair<Value, PTXRegisterMod>> asmValues;
+    PtxBuilder generator(op, rewriter, op.getPtx(), op.hasSideEffect());
+
+    op.getAsmValues(rewriter, asmValues);
+    for (auto &[asmValue, modifier] : asmValues) {
+      LLVM_DEBUG(DBGSNL() << asmValue << "\t Modifier : " << modifier);
+      generator.insertValue(asmValue, modifier);
+    }
+
+    generator.buildAndReplaceOp();
+    return success();
+  }
+};
+
+struct ConvertNVVMToLLVMPass
+    : public impl::ConvertNVVMToLLVMPassBase<ConvertNVVMToLLVMPass> {
+  using Base::Base;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<LLVM::LLVMDialect, NVVM::NVVMDialect>();
+  }
+
+  void runOnOperation() override {
+    ConversionTarget target(getContext());
+    target.addLegalDialect<::mlir::LLVM::LLVMDialect>();
+    RewritePatternSet pattern(&getContext());
+    pattern.add<PtxLowering>(pattern.getContext());
+    if (failed(
+            applyPartialConversion(getOperation(), target, std::move(pattern))))
+      signalPassFailure();
+  }
+};
+
+} // namespace
 No newline at end of file

mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir

+// RUN: mlir-opt --convert-nvvm-to-llvm --split-input-file %s | FileCheck %s
+
+// CHECK-LABEL : @init_mbarrier_arrive_expect_tx
+llvm.func @init_mbarrier_arrive_expect_tx(%barrier : !llvm.ptr<3>, %txcount : i32) -> i32{
+  //CHECK : llvm.inline_asm has_side_effects asm_dialect = att "mbarrier.arrive.expect_tx.shared.b64 %0, [%1], %2;", "=r,r,r" %arg0, %arg1 : (!llvm.ptr<3>, i32) -> i32
+  %res = nvvm.mbarrier.arrive.expect_tx.shared %barrier, %txcount : !llvm.ptr<3>, i32 -> i32
+  llvm.return %res : i32
+}
+
+// CHECK-LABEL : @init_mbarrier_arrive_expect_tx_generic
+llvm.func @init_mbarrier_arrive_expect_tx_generic(%barrier : !llvm.ptr, %txcount : i32)-> i32 {
+  // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "mbarrier.arrive.expect_tx.b64 %0, [%1], %2;", "=r,l,r" %arg0, %arg1 : (!llvm.ptr, i32) -> i32
+  %res = nvvm.mbarrier.arrive.expect_tx %barrier, %txcount : !llvm.ptr, i32 -> i32
+  llvm.return %res : i32
+}
+
+// CHECK-LABEL : @init_mbarrier_try_wait.parity.shared
+llvm.func @init_mbarrier_try_wait.parity.shared(%barrier : !llvm.ptr<3>, %token : i32) -> i32 {
+  // CHECK : llvm.inline_asm has_side_effects asm_dialect = att "{\0A\09.reg .pred P1; \0A\09mbarrier.try_wait.parity.shared.b64 P1, [%1], %2; \0A\09selp.b32 %0, 1, 0, P1; \0A\09}", "=r,r,r" %arg0, %arg1 : (!llvm.ptr<3>, i32) -> i32
+  %res = nvvm.mbarrier.try_wait.parity.shared %barrier, %token : !llvm.ptr<3>, i32 -> i32
+  llvm.return %res : i32
+}
+
+// CHECK-LABEL : @init_mbarrier_try_wait.parity
+llvm.func @init_mbarrier_try_wait.parity(%barrier : !llvm.ptr, %token : i32) -> i32{
+  // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "{\0A\09.reg .pred P1; \0A\09mbarrier.try_wait.parity.b64 P1, [%1], %2; \0A\09selp.b32 %0, 1, 0, P1; \0A\09}", "=r,l,r" %arg0, %arg1 : (!llvm.ptr, i32) -> i32
+  %res = nvvm.mbarrier.try_wait.parity %barrier, %token : !llvm.ptr, i32 -> i32
+  llvm.return %res : i32
+}