diff --git a/lib/parsers/asm-parser-ptx.ts b/lib/parsers/asm-parser-ptx.ts index d1659d367..197147e82 100644 --- a/lib/parsers/asm-parser-ptx.ts +++ b/lib/parsers/asm-parser-ptx.ts @@ -81,7 +81,7 @@ export class PTXAsmParser extends AsmParser { this.labelLine = /^\s*\$?[a-zA-Z_][a-zA-Z0-9_]*:.*$/; - this.hasOpcodeRe = /^\s*(@!?%\w+\s+)?(%[$.A-Z_a-z][\w$.]*\s*=\s*)?[A-Za-z]/; + this.hasOpcodeRe = /^\s*\{?\s*(@!?%\w+\s+)?(%[$.A-Z_a-z][\w$.]*\s*=\s*)?[A-Za-z]/; } override processAsm(asmResult: string, filters: ParseFiltersAndOutputOptions): ParsedAsmResult { @@ -104,6 +104,7 @@ export class PTXAsmParser extends AsmParser { let braceDepth = 0; let lineNumber = 0; let openBraceLineNumber = 0; + let openBraceLineHasOpcode = false; for (let line of asmLines) { const newSource = this.processSourceLine(line, files); @@ -186,10 +187,11 @@ export class PTXAsmParser extends AsmParser { if (this.functionStart.test(line)) { braceDepth++; openBraceLineNumber = lineNumber; + openBraceLineHasOpcode = this.hasOpcode(line); } else if (this.functionEnd.test(line)) { braceDepth--; - // Remove paired braces with no content - if (openBraceLineNumber + 1 === lineNumber) { + // Remove paired braces with no content (but not if the opening brace line had an opcode) + if (openBraceLineNumber + 1 === lineNumber && !openBraceLineHasOpcode) { asm.pop(); continue; } diff --git a/test/asm-parser-tests.ts b/test/asm-parser-tests.ts index b5f96cb3e..436dba34c 100644 --- a/test/asm-parser-tests.ts +++ b/test/asm-parser-tests.ts @@ -72,6 +72,11 @@ describe('PTXAsmParser tests', () => { expect(parser.hasOpcode(' @!%p1 bra LBB6_2;')).toBe(true); expect(parser.hasOpcode(' @%r789 bra LBB6_2;')).toBe(true); }); + it('should identify PTX opcodes wrapped in braces', () => { + expect(parser.hasOpcode('{fma.rn.f16x2 %r9,%r2,%r3,%r270;')).toBe(true); + expect(parser.hasOpcode(' {fma.rn.f16x2 %r9,%r2,%r3,%r270;')).toBe(true); + expect(parser.hasOpcode('{ add.s32 %r1, %r2, %r3;')).toBe(true); + }); }); describe('Nested brace indentation', () => { @@ -158,6 +163,47 @@ mov.u64 %rd1, $str; }); }); + describe('Braced instructions', () => { + it('should not drop FMA instructions wrapped in braces on separate lines', () => { + const input = `{fma.rn.f16x2 %r9,%r2,%r3,%r270; +} +{fma.rn.f16x2 %r13,%r2,%r3,%r9; +}`; + const result = parser.processAsm(input, {}); + const lines = result.asm.map(line => line.text); + + expect(lines.length).toBe(4); + expect(lines[0]).toBe('{fma.rn.f16x2 %r9,%r2,%r3,%r270;'); + expect(lines[1]).toBe('}'); + expect(lines[2]).toBe('{fma.rn.f16x2 %r13,%r2,%r3,%r9;'); + expect(lines[3]).toBe('}'); + }); + + it('should preserve braced instructions inside functions', () => { + const input = `.visible .entry kernel() +{ +mov.u32 %r1, 0; +{fma.rn.f16x2 %r9,%r2,%r3,%r270; +} +{fma.rn.f16x2 %r13,%r2,%r3,%r9; +} +ret; +}`; + const result = parser.processAsm(input, {}); + const lines = result.asm.map(line => line.text); + + expect(lines[0]).toBe('.visible .entry kernel()'); + expect(lines[1]).toBe('{'); + expect(lines[2]).toBe('\tmov.u32 %r1, 0;'); + expect(lines[3]).toBe('\t{fma.rn.f16x2 %r9,%r2,%r3,%r270;'); + expect(lines[4]).toBe('\t}'); + expect(lines[5]).toBe('\t{fma.rn.f16x2 %r13,%r2,%r3,%r9;'); + expect(lines[6]).toBe('\t}'); + expect(lines[7]).toBe('\tret;'); + expect(lines[8]).toBe('}'); + }); + }); + describe('Trim filter', () => { it('should convert tabs to spaces when trim filter is enabled', () => { const input = `\t\t\tcall.uni (retval0),