[mlir][nvvm] Change return type of std::string of getPtx of PtxBuilder

[guraypp]

Summary

getPtx used to return const char*. It is not flexible when one needs to build string in the function. This work changes return type.

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

let description = [{
     If we have following Op definition that returns PTX code by `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;\"; }
+        std::string $cppClass::getPtx() { return std::string(\"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
     ```
InterfaceMethod<
         /*retType=*/"bool",
         /*methodName=*/"hasSideEffect",
         /*args=*/(ins),
         /*methodBody=*/"",
         /*defaultImplementation=*/"return true;"
       >,
     InterfaceMethod<
         /*desc=*/[{ Returns PTX code. }],
-        /*retType=*/"const char*",
+        /*retType=*/"std::string",
         /*methodName=*/"getPtx"
       >,
     InterfaceMethod<
         /*desc=*/[{Generate constant value.}],
         /*retType=*/"::mlir::Value",
         /*methodName=*/"makeConstantI32",
         /*args=*/(ins "::mlir::RewriterBase &":$rewriter, "unsigned" : $val),
         /*methodBody=*/"",

 }
 
 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;"; }
+    std::string $cppClass::getPtx() { return std::string("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;"; }
+    std::string $cppClass::getPtx() { return std::string("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"
+    std::string $cppClass::getPtx() {
+      return std::string("{\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"
+    std::string $cppClass::getPtx() {
+      return std::string("{\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:$state)> {
   string llvmBuilder = [{
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({makeConstantI32(rewriter, getSize()), PTXRegisterMod::Read});
       asmValues.push_back({getCpSize(), PTXRegisterMod::Read});
     }        
   }];
   let extraClassDefinition = [{        
-    const char* $cppClass::getPtx() { 
+    std::string $cppClass::getPtx() { 
       if(getModifier() == NVVM::LoadCacheModifierKind::CG)
-        return "cp.async.cg.shared.global [%0], [%1], %2, %3;\n"; 
+        return std::string("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";        
+        return std::string("cp.async.ca.shared.global [%0], [%1], %2, %3;\n");
       llvm_unreachable("unsupported cache modifier");      
     }
   }];
 }
 
 def NVVM_CpAsyncCommitGroupOp : NVVM_Op<"cp.async.commit.group"> {
   string llvmBuilder = [{
       createIntrinsicCall(builder, llvm::Intrinsic::nvvm_cp_async_commit_group);

mlir/lib/Conversion/NVVMToLLVM/NVVMToLLVM.cpp

 #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 <regex>
 #include <string>
 
 #define DEBUG_TYPE "nvvm-to-llvm"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
 #define DBGSNL() (llvm::dbgs() << "\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;
+  std::string 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>())
if (auto ptr = v.getType().dyn_cast<LLVM::LLVMPointerType>()) {
       }
       return 'l';
     }
     assert(false && "Register type is not handled yet");
     return ' ';
   }
 
 public:
-  PtxBuilder(Operation *op, PatternRewriter &rewriter, const char *ptxAsm,
+  PtxBuilder(Operation *op, PatternRewriter &rewriter, std::string ptxAsm,
              bool sideEffects = false)
-      : op(op), rewriter(rewriter), asmStr(ptxAsm), sideEffects(sideEffects) {}
+      : op(op), rewriter(rewriter), asmStr(std::move(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 << "+";
LLVM::InlineAsmOp build() {
     Type resultType = hasResult ? op->getResult(0).getType()
                                 : LLVM::LLVMVoidType::get(op->getContext());
 
     // Remove the last comma from the constraints string.
     if (!asmConstraints.empty() &&
         asmConstraints[asmConstraints.size() - 1] == ',')
       asmConstraints.pop_back();
 
+    // asm keywords expects %, but inline assembly uses $. Replace all % with $
+    std::replace(asmStr.begin(), asmStr.end(), '%', '$');
+
     return rewriter.create<LLVM::InlineAsmOp>(
         op->getLoc(), resultType,
         /*operands=*/asmVals,
-        /*asm_string=*/asmStr,
+        /*asm_string=*/llvm::StringRef(asmStr),
         /*constraints=*/asmConstraints.data(),
         /*has_side_effects=*/sideEffects,
         /*is_align_stack=*/false,
         /*asm_dialect=*/asmDialectAttr,
         /*operand_attrs=*/ArrayAttr());
   }
 
   void buildAndReplaceOp() {

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" %{{.*}}, %{{.*}} : (!llvm.ptr<3>, i32) -> i32            
-  %res = nvvm.mbarrier.arrive.expect_tx.shared %barrier, %txcount : !llvm.ptr<3>, i32 -> i32
-  llvm.return %res : i32
+llvm.func @init_mbarrier_arrive_expect_tx(%barrier : !llvm.ptr<3>, %txcount : i32) -> i64 {
+  //CHECK : llvm.inline_asm has_side_effects asm_dialect = att "mbarrier.arrive.expect_tx.shared.b64 $0, [$1], $2;", "=l,r,r" %{{.*}}, %{{.*}} : (!llvm.ptr<3>, i32) -> i64            
+  %res = nvvm.mbarrier.arrive.expect_tx.shared %barrier, %txcount : !llvm.ptr<3>, i32 -> i64
+  llvm.return %res : i64
 }
 
 // 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" %{{.*}}, %{{.*}} : (!llvm.ptr, i32) -> i32
-  %res = nvvm.mbarrier.arrive.expect_tx %barrier, %txcount : !llvm.ptr, i32 -> i32
-  llvm.return %res : i32
+llvm.func @init_mbarrier_arrive_expect_tx_generic(%barrier : !llvm.ptr, %txcount : i32)-> i64 {
+  // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "mbarrier.arrive.expect_tx.b64 $0, [$1], $2;", "=l,l,r" %{{.*}}, %{{.*}} : (!llvm.ptr, i32) -> i64
+  %res = nvvm.mbarrier.arrive.expect_tx %barrier, %txcount : !llvm.ptr, i32 -> i64
+  llvm.return %res : i64
 }
 
 // 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" %{{.*}}, %{{.*}} : (!llvm.ptr<3>, i32) -> i32
+llvm.func @init_mbarrier_try_wait_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" %{{.*}}, %{{.*}} : (!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" %{{.*}}, %{{.*}} : (!llvm.ptr, i32) -> i32
+llvm.func @init_mbarrier_try_wait(%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" %{{.*}}, %{{.*}} : (!llvm.ptr, i32) -> i32
   %res = nvvm.mbarrier.try_wait.parity %barrier, %token : !llvm.ptr, i32 -> i32
   llvm.return %res : i32
 }
 
 // CHECK-LABEL : @async_cp
 func.func @async_cp(%dst: !llvm.ptr<3>, %src: !llvm.ptr<1>) {
   // CHECK : nvvm.cp.async.shared.global %{{.*}}, %{{.*}}, 16, cache =  ca : !llvm.ptr<3>, !llvm.ptr<1>
   nvvm.cp.async.shared.global %dst, %src, 16, cache =  ca : !llvm.ptr<3>, !llvm.ptr<1>
   // CHECK : nvvm.cp.async.shared.global %{{.*}}, %{{.*}}, 16, cache =  cg : !llvm.ptr<3>, !llvm.ptr<1>
   nvvm.cp.async.shared.global %dst, %src, 16, cache =  cg : !llvm.ptr<3>, !llvm.ptr<1>
   return
 }
 
 // CHECK-LABEL : @async_cp_zfill
 func.func @async_cp_zfill(%dst: !llvm.ptr<3>, %src: !llvm.ptr<1>, %cpSize: i32) {
-  // CHECK : llvm.inline_asm has_side_effects asm_dialect = att "cp.async.cg.shared.global [%0], [%1], %2, %3;\0A", "r,l,r" %{{.*}}, %{{.*}}, %{{.*}} : (!llvm.ptr<3>, !llvm.ptr<1>, i32) -> !llvm.void
+  // CHECK : llvm.inline_asm has_side_effects asm_dialect = att "cp.async.cg.shared.global [$0], [$1], $2, $3;\0A", "r,l,r" %{{.*}}, %{{.*}}, %{{.*}} : (!llvm.ptr<3>, !llvm.ptr<1>, i32) -> !llvm.void
   nvvm.cp.async.shared.global %dst, %src, 16, cache =  cg, %cpSize : !llvm.ptr<3>, !llvm.ptr<1>, i32
-  // CHECK : llvm.inline_asm has_side_effects asm_dialect = att "cp.async.ca.shared.global [%0], [%1], %2, %3;\0A", "r,l,r" %{{.*}}, %{{.*}}, %{{.*}} : (!llvm.ptr<3>, !llvm.ptr<1>, i32) -> !llvm.void
+  // CHECK : llvm.inline_asm has_side_effects asm_dialect = att "cp.async.ca.shared.global [$0], [$1], $2, $3;\0A", "r,l,r" %{{.*}}, %{{.*}}, %{{.*}} : (!llvm.ptr<3>, !llvm.ptr<1>, i32) -> !llvm.void
   nvvm.cp.async.shared.global %dst, %src, 4, cache =  ca, %cpSize : !llvm.ptr<3>, !llvm.ptr<1>, i32
   return
 }