1 files changed, 41 insertions, 1 deletions
diff --git a/llvm/utils/UpdateTestChecks/asm.py b/llvm/utils/UpdateTestChecks/asm.py
index bb59484..01a25a6 100644
--- a/llvm/utils/UpdateTestChecks/asm.py
+++ b/llvm/utils/UpdateTestChecks/asm.py
@@ -178,6 +178,31 @@ ASM_FUNCTION_CSKY_RE = re.compile(
     r'.Lfunc_end[0-9]+:\n',
     flags=(re.M | re.S))
 
+ASM_FUNCTION_NVPTX_RE = re.compile(
+    # function attributes and retval
+    # .visible .func (.param .align 16 .b8 func_retval0[32])
+    #r'^(\.visible\s+)?\.func\s+(\([^\)]*\)\s*)?'
+    r'^(\.(func|visible|weak|entry|noreturn|extern)\s+)+(\([^\)]*\)\s*)?'
+
+    # function name
+    r'(?P<func>[^\(\n]+)'
+
+    # function name separator (opening brace)
+    r'(?P<func_name_separator>\()'
+
+    # function parameters
+    # (
+    #   .param .align 16 .b8 callee_St8x4_param_0[32]
+    # ) // -- Begin function callee_St8x4
+    r'[^\)]*\)(\s*//[^\n]*)?\n'
+
+    # function body
+    r'(?P<body>.*?)\n'
+
+    # function body end marker
+    r'\s*// -- End function',
+    flags=(re.M | re.S))
+
 SCRUB_X86_SHUFFLES_RE = (
     re.compile(
         r'^(\s*\w+) [^#\n]+#+ ((?:[xyz]mm\d+|mem)( \{%k\d+\}( \{z\})?)? = .*)$',
@@ -388,6 +413,20 @@ def scrub_asm_csky(asm, args):
   asm = common.SCRUB_TRAILING_WHITESPACE_RE.sub(r'', asm)
   return asm
 
+def scrub_asm_nvptx(asm, args):
+  # Scrub runs of whitespace out of the assembly, but leave the leading
+  # whitespace in place.
+  asm = common.SCRUB_WHITESPACE_RE.sub(r' ', asm)
+  # Expand the tabs used for indentation.
+  asm = string.expandtabs(asm, 2)
+  # Strip trailing whitespace.
+  asm = common.SCRUB_TRAILING_WHITESPACE_RE.sub(r'', asm)
+  return asm
+
+
+# Returns a tuple of a scrub function and a function regex. Scrub function is
+# used to alter function body in some way, for example, remove traling spaces.
+# Function regex is used to match function name, body, etc. in raw llc output.
 def get_run_handler(triple):
   target_handlers = {
       'i686': (scrub_asm_x86, ASM_FUNCTION_X86_RE),
@@ -426,6 +465,7 @@ def get_run_handler(triple):
       'wasm32': (scrub_asm_wasm32, ASM_FUNCTION_WASM32_RE),
       've': (scrub_asm_ve, ASM_FUNCTION_VE_RE),
       'csky': (scrub_asm_csky, ASM_FUNCTION_CSKY_RE),
+      'nvptx': (scrub_asm_nvptx, ASM_FUNCTION_NVPTX_RE)
   }
   handler = None
   best_prefix = ''
@@ -444,7 +484,7 @@ def get_run_handler(triple):
 def add_checks(output_lines, comment_marker, prefix_list, func_dict,
                    func_name, is_filtered):
   # Label format is based on ASM string.
-  check_label_format = '{} %s-LABEL: %s%s:'.format(comment_marker)
+  check_label_format = '{} %s-LABEL: %s%s%s'.format(comment_marker)
   global_vars_seen_dict = {}
   common.add_checks(output_lines, comment_marker, prefix_list, func_dict,
                     func_name, check_label_format, True, False,