diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/target/detect.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/target/detect.rs index 20ce82e87a..2bfff0af70 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/target/detect.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/target/detect.rs @@ -5,6 +5,7 @@ use super::features::{DetectedFeatures, ModuleRequirements, PtxIsaRequirement}; use crate::error::PipelineError; +use std::borrow::Cow; use std::path::Path; pub(super) fn contains_wgmma_features(contents: &str) -> bool { @@ -95,11 +96,266 @@ fn contains_multimem_blackwell_features(contents: &str) -> bool { }) } -/// Checks the PTX 8.6 floating-point extension to `redux.sync`. +/// Checks the PTX 8.6 floating-point extension to `redux.sync`: a call to a +/// float NVVM redux intrinsic, or a `redux.sync ... .f32` instruction in +/// inline PTX. Decided on tokens, so text that only mentions either -- in a +/// comment, a string or another symbol's name -- does not count. fn contains_redux_f32_features(contents: &str) -> bool { - contents - .split(';') - .any(|statement| statement.contains("redux.sync") && statement.contains(".f32")) + let tokens = lex_llvm_ir(contents); + tokens.iter().enumerate().any(|(index, token)| match token { + IrToken::Global(name) => is_float_redux_callee(name), + IrToken::Word(word) if *word == "asm" => inline_asm_template(&tokens[index + 1..]) + .is_some_and(|template| ptx_contains_float_redux(&unescape_llvm_string(template))), + // A raw PTX fragment rather than IR, which is how tests pass PTX. + IrToken::Word(word) => ptx_contains_float_redux(word), + IrToken::Label(_) | IrToken::String(_) => false, + }) +} + +/// `llvm.nvvm.redux.sync.` followed by a float operation. Those are the ones +/// beginning with `f` (`fmin`, `fmax`, each with optional `.abs` and `.NaN`); +/// no integer operation does. +fn is_float_redux_callee(name: &str) -> bool { + name.strip_prefix("llvm.nvvm.redux.sync.") + .is_some_and(|operation| operation.starts_with('f')) +} + +/// Does this PTX contain a `redux` instruction with an `.f32` modifier? +/// +/// An opcode's modifiers are the run of `.name` tokens after it, in any +/// order, so `redux.sync.min.f32`, `redux.f32.sync.min` and +/// `redux .sync .min .f32` are the same instruction. ptxas requires `.sync` +/// on every `redux`, so it is not checked. +fn ptx_contains_float_redux(ptx: &str) -> bool { + if !ptx.contains("redux") { + return false; + } + let tokens = lex_ptx(ptx); + tokens.iter().enumerate().any(|(index, token)| { + *token == PtxToken::Identifier("redux") + && tokens[index + 1..] + .iter() + .map_while(|token| match token { + PtxToken::Modifier(modifier) => Some(*modifier), + _ => None, + }) + .any(|modifier| modifier == ".f32") + }) +} + +/// A PTX token. Whitespace, comments and string literals produce none. +#[derive(Debug, PartialEq)] +enum PtxToken<'a> { + /// An opcode or other identifier: `redux`, `%r1`, `$0`. + Identifier(&'a str), + /// A `.name` token: an instruction modifier or a directive. + Modifier(&'a str), + /// Punctuation or a number. + Other, +} + +/// Split PTX into [`PtxToken`]s, dropping whitespace, `//` and `/* */` +/// comments, and string literals such as a `.file` path. +fn lex_ptx(ptx: &str) -> Vec> { + fn is_identifier(byte: u8) -> bool { + byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'$') + } + + let bytes = ptx.as_bytes(); + // In an asm template a bare `$` starts an operand and `$$` is a literal + // dollar, so `.f32$0` is the modifier `.f32` followed by `$0`. + let run_end = |from: usize| { + let mut end = from; + while end < bytes.len() && is_identifier(bytes[end]) { + if bytes[end] == b'$' { + if bytes.get(end + 1) != Some(&b'$') { + break; + } + end += 1; + } + end += 1; + } + end + }; + + let mut tokens = Vec::new(); + let mut index = 0; + while index < bytes.len() { + let rest = &bytes[index..]; + if rest.starts_with(b"//") { + index += rest + .iter() + .position(|&byte| byte == b'\n') + .unwrap_or(rest.len()); + } else if rest.starts_with(b"/*") { + index += rest[2..] + .windows(2) + .position(|pair| pair == b"*/") + .map_or(rest.len(), |offset| offset + 4); + } else if rest[0] == b'"' { + // A backslash is literal in a PTX string -- ptxas accepts + // `.file 1 "C:\dir\"` -- so the next quote always closes it. + index += rest[1..] + .iter() + .position(|&byte| byte == b'"') + .map_or(rest.len(), |offset| offset + 2); + } else if rest[0] == b'.' && rest.get(1).is_some_and(|&byte| is_identifier(byte)) { + let end = run_end(index + 1); + tokens.push(PtxToken::Modifier(&ptx[index..end])); + index = end; + } else if rest[0].is_ascii_digit() { + // A number consumes its whole run -- `0f3F800000`, `0x1F`, `123U` -- + // so an identifier is never read out of its tail. + index = run_end(index + 1); + tokens.push(PtxToken::Other); + } else if rest[0].is_ascii_alphabetic() || matches!(rest[0], b'_' | b'$' | b'%') { + let end = run_end(index + 1); + tokens.push(PtxToken::Identifier(&ptx[index..end])); + index = end; + } else { + if !rest[0].is_ascii_whitespace() { + tokens.push(PtxToken::Other); + } + index += 1; + } + } + tokens +} + +/// The template of an inline-asm expression, given the tokens after `asm`: +/// the string after `[sideeffect] [alignstack] [inteldialect] [unwind]`, +/// each optional and in that order, as `LLParser` accepts them. Also covers +/// `module asm`. +fn inline_asm_template<'a>(after_asm: &[IrToken<'a>]) -> Option<&'a str> { + const MODIFIERS: [&str; 4] = ["sideeffect", "alignstack", "inteldialect", "unwind"]; + let mut tokens = after_asm.iter(); + let mut next = tokens.next(); + for modifier in MODIFIERS { + if matches!(next, Some(IrToken::Word(word)) if *word == modifier) { + next = tokens.next(); + } + } + match next { + Some(IrToken::String(template)) => Some(template), + _ => None, + } +} + +/// An LLVM IR token. Comments, punctuation, local and metadata names and +/// attribute groups produce none. +#[derive(Debug, PartialEq)] +enum IrToken<'a> { + /// A bare keyword, type or identifier. + Word(&'a str), + /// A basic-block label, `name:`. + Label(&'a str), + /// A global symbol with the sigil removed and quoting decoded, so + /// `@name` and `@"quoted name"` both yield the name. + Global(Cow<'a, str>), + /// A string literal between its quotes, still escaped. + String(&'a str), +} + +/// Split LLVM IR into [`IrToken`]s. Quoted text is consumed whole before +/// anything inside it is read, so a `;` in a string is not a comment and a +/// `@` in a quoted name does not start another symbol. +fn lex_llvm_ir(contents: &str) -> Vec> { + fn is_identifier(byte: u8) -> bool { + byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'$' | b'.' | b'_') + } + + let bytes = contents.as_bytes(); + // LLVM writes a quote inside quoted text as `\22`, so the next `"` closes it. + let closing_quote = |from: usize| { + bytes[from..] + .iter() + .position(|&byte| byte == b'"') + .map_or(bytes.len(), |offset| from + offset) + }; + let identifier_end = |from: usize| { + bytes[from..] + .iter() + .position(|&byte| !is_identifier(byte)) + .map_or(bytes.len(), |offset| from + offset) + }; + + let mut tokens = Vec::new(); + let mut index = 0; + while index < bytes.len() { + match bytes[index] { + b';' => { + index = bytes[index..] + .iter() + .position(|&byte| byte == b'\n') + .map_or(bytes.len(), |offset| index + offset); + } + b'"' => { + let close = closing_quote(index + 1); + tokens.push(IrToken::String(&contents[index + 1..close])); + index = bytes.len().min(close + 1); + } + sigil @ (b'@' | b'%' | b'!' | b'#') => { + let (name, next) = if bytes.get(index + 1) == Some(&b'"') { + let close = closing_quote(index + 2); + ( + unescape_llvm_string(&contents[index + 2..close]), + bytes.len().min(close + 1), + ) + } else { + let end = identifier_end(index + 1); + (Cow::Borrowed(&contents[index + 1..end]), end) + }; + if sigil == b'@' { + tokens.push(IrToken::Global(name)); + } + index = next; + } + byte if is_identifier(byte) => { + let end = identifier_end(index); + let word = &contents[index..end]; + if bytes.get(end) == Some(&b':') { + tokens.push(IrToken::Label(word)); + index = end + 1; + } else { + tokens.push(IrToken::Word(word)); + index = end; + } + } + _ => index += 1, + } + } + tokens +} + +/// Decode LLVM quoted text as `llvm::UnEscapeLexed` does: `\\` is one +/// backslash, `\XX` is the byte with that hex value, and any other backslash +/// stands for itself. +fn unescape_llvm_string(text: &str) -> Cow<'_, str> { + if !text.contains('\\') { + return Cow::Borrowed(text); + } + let hex = |byte: Option<&u8>| byte.and_then(|&byte| (byte as char).to_digit(16)); + let bytes = text.as_bytes(); + let mut decoded = Vec::with_capacity(bytes.len()); + let mut index = 0; + while index < bytes.len() { + if bytes[index] == b'\\' { + if bytes.get(index + 1) == Some(&b'\\') { + decoded.push(b'\\'); + index += 2; + continue; + } + if let (Some(high), Some(low)) = (hex(bytes.get(index + 1)), hex(bytes.get(index + 2))) + { + decoded.push((high * 16 + low) as u8); + index += 3; + continue; + } + } + decoded.push(bytes[index]); + index += 1; + } + Cow::Owned(String::from_utf8_lossy(&decoded).into_owned()) } /// Checks for forward-compatible instructions whose minimum target is sm_90. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/target/tests/detection_async_memory.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/target/tests/detection_async_memory.rs index 9c644810bb..9f24089ec6 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/target/tests/detection_async_memory.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/target/tests/detection_async_memory.rs @@ -586,3 +586,209 @@ fn test_tma_multicast_detection_requires_cta_mask() { ); assert!(!contains_tma_cta_group_features(unrelated_i32)); } + +/// `redux.sync` is f32 by its own token, never by an `.f32` sitting nearby +/// (#1303). +#[test] +fn test_redux_f32_detection_is_scoped_to_the_instruction() { + // The shape that found this: integer reductions in one kernel and + // ordinary float work in another. In the module that was rejected, the + // two matches were 547 lines apart. + let integer_redux_and_distant_f32 = r#" + declare i32 @llvm.nvvm.redux.sync.add(i32, i32) #0 + declare float @llvm.nvvm.shfl.sync.bfly.f32(i32, float, i32, i32) #0 + "#; + // LLVM does not require a newline between declarations, so scoping to a + // line is not enough either. + let same_line = "declare i32 @llvm.nvvm.redux.sync.add(i32, i32) declare float @llvm.nvvm.shfl.sync.bfly.f32(i32, float, i32, i32)"; + // An SSA name may end in `.f32` and says nothing about the callee. + let unrelated_ssa_name = "%acc.f32 = call i32 @llvm.nvvm.redux.sync.umax(i32 -1, i32 %v)"; + + for text in [integer_redux_and_distant_f32, same_line, unrelated_ssa_name] { + assert!( + !detect_features_in_llvm_text(text).contains(DetectedFeatures::ReduxF32), + "an integer reduction must not select the f32 extension: {text}" + ); + } + + // The intrinsic spelling has no type suffix -- the `f`-prefixed operation + // is what makes it float. These were detected before only when some + // unrelated `.f32` happened to sit nearby. + let float_intrinsic = "%r = call float @llvm.nvvm.redux.sync.fmin(float %v, i32 -1)"; + let float_intrinsic_qualified = + "%r = call float @llvm.nvvm.redux.sync.fmax.abs.NaN(float %v, i32 -1)"; + // Inline PTX carries the type as a modifier of the opcode token. + let float_asm = r#"call float asm sideeffect "redux.sync.min.abs.NaN.f32 $0, $1, $2;", "=f,f,r"(float %v, i32 %m)"#; + + for text in [float_intrinsic, float_intrinsic_qualified, float_asm] { + assert!( + detect_features_in_llvm_text(text).contains(DetectedFeatures::ReduxF32), + "a float reduction must select the f32 extension: {text}" + ); + } +} + +/// Mentioning an intrinsic is not using one, and a real call may be spelled +/// with a quoted identifier (#1303). +#[test] +fn test_redux_f32_detection_reads_code_and_not_prose() { + // Text that names the callee without calling it. + let in_a_comment = + "%r = call i32 @llvm.nvvm.redux.sync.add(i32 -1, i32 %v) ; @llvm.nvvm.redux.sync.fmin"; + let in_a_data_string = r#"@.msg = private constant [27 x i8] c"@llvm.nvvm.redux.sync.fmin\00""#; + let in_metadata = r#"!0 = !{!"@llvm.nvvm.redux.sync.fmin"}"#; + // The same for the inline-PTX spelling: a comment and a data constant are + // not instructions. + let ptx_in_a_comment = " ; redux.sync.min.abs.NaN.f32 $0, $1, $2;"; + let ptx_in_a_data_string = r#"@.msg = private constant [25 x i8] c"redux.sync.min.f32 a,b,c;""#; + + for text in [ + in_a_comment, + in_a_data_string, + in_metadata, + ptx_in_a_comment, + ptx_in_a_data_string, + ] { + assert!( + !detect_features_in_llvm_text(text).contains(DetectedFeatures::ReduxF32), + "text that only mentions the f32 reduction must not select it: {text}" + ); + } + + // A real call, spelled with a quoted identifier and with one of its + // characters escaped -- `\66` is `f`, so both name `fmin`. + let quoted = r#"%r = call float @"llvm.nvvm.redux.sync.fmin"(float %v, i32 -1)"#; + let quoted_escaped = r#"%r = call float @"llvm.nvvm.redux.sync.\66min"(float %v, i32 -1)"#; + // And the constraint string beside an asm template is data, while the + // template itself is code. + let asm_template = r#"call float asm sideeffect "redux.sync.min.abs.NaN.f32 $0, $1, $2;", "=f,f,r"(float %v, i32 %m)"#; + + for text in [quoted, quoted_escaped, asm_template] { + assert!( + detect_features_in_llvm_text(text).contains(DetectedFeatures::ReduxF32), + "a float reduction must select the f32 extension: {text}" + ); + } +} + +/// Token identity at the edges the flattened search lost (#1303). +/// +/// Each input is a line from a module `llvm-as` from the pinned toolchain +/// accepts, so none of these is a shape valid IR cannot produce. The escape +/// cases were also round-tripped through `llvm-dis` to confirm what LLVM +/// itself decodes them to. +#[test] +fn test_redux_f32_detection_keeps_token_identity() { + // `unwind` is the last of the four modifiers `asm` may carry, in order: + // sideeffect, alignstack, inteldialect, unwind. + let unwind = r#"%r = call float asm sideeffect unwind "redux.sync.min.f32 $0, $1, $2;", "=f,f,r"(float %v, i32 %m)"#; + // `\72` is `r`: LLVM decodes the template before the backend sees it. + let escaped_opcode = r#"%r = call float asm sideeffect "\72edux.sync.min.f32 $0, $1, $2;", "=f,f,r"(float %v, i32 %m)"#; + // A PTX comment ended by an *escaped* newline, then a real instruction. + // Only decoding before stripping comments finds it: the raw template is + // one line, so a `//` there would appear to run to its end. + let after_escaped_newline = r#"%r = call float asm sideeffect "// scale first\0Aredux.sync.min.f32 $0, $1, $2;", "=f,f,r"(float %v, i32 %m)"#; + // Module-level inline asm has the same grammar and reaches the same PTX. + let module_asm = r#"module asm "redux.sync.min.f32 %f1, %f2, %r3;""#; + + for text in [unwind, escaped_opcode, after_escaped_newline, module_asm] { + assert!( + detect_features_in_llvm_text(text).contains(DetectedFeatures::ReduxF32), + "a float reduction must select the f32 extension: {text}" + ); + } + + // PTX comments inside a template are prose, in either comment form. + let ptx_line_comment = r#"%r = call float asm sideeffect "// redux.sync.min.f32 $0, $1, $2;\0Amov.f32 $0, $1;", "=f,f,r"(float %v, i32 %m)"#; + let ptx_block_comment = r#"%r = call float asm sideeffect "/* redux.sync.min.f32 $0, $1, $2; */ mov.f32 $0, $1;", "=f,f,r"(float %v, i32 %m)"#; + // `\\` is one backslash and the hex digits after it are plain text: `llc` + // emits this template as `\20redux.sync.min.f32 ...`, which is not an + // instruction. A decoder that skipped the `\\` rule would read `\20` as a + // space instead and find a `redux.sync.min.f32` that is not there. + let escaped_backslash = r#"%r = call float asm sideeffect "\\20redux.sync.min.f32 $0, $1, $2;", "=f,f,r"(float %v, i32 %m)"#; + // One symbol whose quoted name happens to contain another's spelling. + let embedded_in_a_name = r#"@"unused @llvm.nvvm.redux.sync.fmin" = global i32 0"#; + // A basic block may be named like an instruction without being one. + let label = "redux.sync.min.f32:\n ret void"; + + for text in [ + ptx_line_comment, + ptx_block_comment, + escaped_backslash, + embedded_in_a_name, + label, + ] { + assert!( + !detect_features_in_llvm_text(text).contains(DetectedFeatures::ReduxF32), + "no float reduction is executed here: {text}" + ); + } +} + +/// PTX lexical rules inside an asm template (#1303). +/// +/// Every input is a complete module accepted by `llvm-as` and `llc` from the +/// pinned toolchain and by ptxas at sm_100a. The float cases are the ones +/// ptxas then refuses at sm_80 as `redux.f32`; the others assemble there. +#[test] +fn test_redux_f32_detection_reads_ptx_as_ptx() { + // Whitespace, a newline or a comment may separate an opcode from its + // modifiers; it is still one instruction. + let spaced_type = r#"%r = call float asm sideeffect "redux.sync.min .f32 $0, $1, $2;", "=f,f,r"(float %v, i32 %m)"#; + let comment_between = r#"%r = call float asm sideeffect "redux.sync.min/* qualifier */.f32 $0, $1, $2;", "=f,f,r"(float %v, i32 %m)"#; + let fully_spaced = r#"%r = call float asm sideeffect "redux .sync .min .f32 $0, $1, $2;", "=f,f,r"(float %v, i32 %m)"#; + let newline_between = r#"%r = call float asm sideeffect "redux.sync.min\0A.f32 $0, $1, $2;", "=f,f,r"(float %v, i32 %m)"#; + // A backslash is literal inside a PTX string, so `"nounroll\"` closes at + // its last quote and the instruction after it is real. + let after_trailing_backslash = r#"%r = call float asm sideeffect ".pragma \22nounroll\5C\22;\0Aredux.sync.min.f32 $0, $1, $2;", "=f,f,r"(float %v, i32 %m)"#; + // A bare `$` starts an operand, so it ends the modifier before it. llc + // emits `redux.sync.min.f32%r1, %r2, %r3;`, which ptxas assembles at + // sm_100a and refuses at sm_80 with `Instruction 'redux.f32' not supported`. + let glued_operand = r#"%r = call float asm sideeffect "redux.sync.min.f32$0, $1, $2;", "=f,f,r"(float %v, i32 %m)"#; + // Modifiers may come in any order; ptxas gives both of these the same + // two results as the glued form. + let sync_after_operation = r#"%r = call float asm sideeffect "redux.min.sync.f32 $0, $1, $2;", "=f,f,r"(float %v, i32 %m)"#; + let type_first = r#"%r = call float asm sideeffect "redux.f32.sync.min $0, $1, $2;", "=f,f,r"(float %v, i32 %m)"#; + + for text in [ + spaced_type, + comment_between, + fully_spaced, + newline_between, + after_trailing_backslash, + glued_operand, + sync_after_operation, + type_first, + ] { + assert!( + detect_features_in_llvm_text(text).contains(DetectedFeatures::ReduxF32), + "a float reduction must select the f32 extension: {text}" + ); + } + + // A directive's quoted data is not an instruction. + let file_path = r#"module asm ".file 1 \22name redux.sync.min.f32 label.cu\22""#; + // Spacing does not make an integer reduction float. + let spaced_integer = r#"%r = call i32 asm sideeffect "redux .sync .min .u32 $0, $1, $2;", "=r,r,r"(i32 %v, i32 %m)"#; + // Nor does float work beside it in the same template: `.f32` counts only + // as a modifier of the `redux` opcode itself. + let integer_beside_f32 = r#"%r = call i32 asm sideeffect "{ .reg .f32 t; redux.sync.min.u32 $0, $1, $2; mov.f32 t, 0f3F800000; }", "=r,r,r"(i32 %v, i32 %m)"#; + // The integer forms of the glued and reordered inputs above; ptxas + // assembles both at sm_80. + let glued_integer = + r#"%r = call i32 asm sideeffect "redux.sync.min.u32$0, $1, $2;", "=r,r,r"(i32 %v, i32 %m)"#; + let reordered_integer = r#"%r = call i32 asm sideeffect "redux.min.sync.u32 $0, $1, $2;", "=r,r,r"(i32 %v, i32 %m)"#; + + for text in [ + file_path, + spaced_integer, + integer_beside_f32, + glued_integer, + reordered_integer, + ] { + assert!( + !detect_features_in_llvm_text(text).contains(DetectedFeatures::ReduxF32), + "no float reduction is executed here: {text}" + ); + } +}