Fix decoration parsing in spirq-spvasm

spirq-as/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "spirq-as"
-version = "0.1.2"
+version = "0.1.3"
 authors = ["PENGUINLIONG <[email protected]>"]
 edition = "2018"
 license = "MIT OR Apache-2.0"
@@ -19,5 +19,5 @@ maintenance = { status = "actively-developed" }
 
 [dependencies]
 anyhow = "1.0"
-spirq-spvasm = { version = "0.1.2", path = "../spirq-spvasm" }
+spirq-spvasm = { version = "0.1.3", path = "../spirq-spvasm" }
 clap = { version = "4.0.6", features = ["derive"] }

spirq-dis/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "spirq-dis"
-version = "0.1.2"
+version = "0.1.3"
 authors = ["PENGUINLIONG <[email protected]>"]
 edition = "2018"
 license = "MIT OR Apache-2.0"
@@ -19,5 +19,5 @@ maintenance = { status = "actively-developed" }
 
 [dependencies]
 anyhow = "1.0"
-spirq-spvasm = { version = "0.1.2", path = "../spirq-spvasm" }
+spirq-spvasm = { version = "0.1.3", path = "../spirq-spvasm" }
 clap = { version = "4.0.6", features = ["derive"] }

spirq-spvasm/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "spirq-spvasm"
-version = "0.1.2"
+version = "0.1.3"
 authors = ["PENGUINLIONG <[email protected]>"]
 edition = "2018"
 license = "MIT OR Apache-2.0"

spirq-spvasm/scripts/generate_decorate_parameter_enum_type.py

@@ -0,0 +1,77 @@
+import json
+
+with open("assets/spirv/spirv.core.grammar.json") as f:
+    j = json.load(f)
+
+j = [
+    parameter_kind
+    for parameter_kind in j["operand_kinds"]
+    if parameter_kind["kind"] == "Decoration"
+][0]
+
+parameter_kinds = {}
+for instr in j["enumerants"]:
+    value = instr["value"]
+
+    op_parameter_kinds = {}
+    if "parameters" in instr:
+        parameters = instr["parameters"]
+        i = 0
+        for parameter in parameters:
+            kind = parameter["kind"]
+            if kind in ["IdResultType", "IdResult"]:
+                continue
+            if (
+                kind.startswith("Id")
+                or kind.startswith("Literal")
+                or kind.startswith("Pair")
+            ):
+                i += 1
+                continue
+
+            op_parameter_kinds[i] = parameter["kind"]
+            i += 1
+    parameter_kinds[value] = op_parameter_kinds
+
+out = []
+
+out += [
+    "use anyhow::{bail, Result};",
+    "use num_traits::FromPrimitive;",
+    "use spirq_core::spirv::Op;",
+    "",
+    "fn unknown_decorate_parameter_index(decoration: u32, i: usize) -> Result<&'static str> {",
+    '    let opname = Op::from_u32(decoration).map(|op| format!("{:?}", op)).unwrap_or("<unknown>".to_owned());',
+    '    bail!("Unknown op {} ({}) parameter index: {}", opname, decoration, i)',
+    "}",
+    "",
+    "pub fn decorate_parameter_enum_type(decoration: u32, i: usize) -> Result<&'static str> {",
+    "    let out: &'static str = match decoration {",
+]
+
+for decoration, decorate_parameter_kinds in parameter_kinds.items():
+    if len(decorate_parameter_kinds) == 0:
+        continue
+
+    out += [
+        f"        {decoration} => match i {{",
+    ]
+    for i, kind in decorate_parameter_kinds.items():
+        out += [
+            f'            {i} => "{kind}",',
+        ]
+    out += [
+        "            _ => return unknown_decorate_parameter_index(decoration, i),",
+        "        },",
+    ]
+
+out += [
+    '        _ => bail!("{}-th parameter of decoration {} is not a enum", i, decoration),',
+    "    };",
+    "    Ok(out)",
+    "}",
+    "",
+]
+
+with open("spirq-spvasm/src/generated/decorate_parameter_enum_type.rs", "w") as f:
+    f.write("\n".join(out))

spirq-spvasm/src/asm/assembler.rs

@@ -484,17 +484,96 @@ impl Assembler {
             .build();
         Ok(instr.into_words())
     }
+    fn assemble_op_decorate(&mut self, instr: &Instruction) -> Result<Vec<u32>> {
+        if instr.operands.len() < 2 {
+            bail!("OpDecorate expected at least 2 operands");
+        }
+        let target_id = match &instr.operands[0] {
+            Operand::IdRef(IdRef::Id(id)) => *id,
+            _ => bail!("OpDecorate expected target id"),
+        };
+        let decoration = match &instr.operands[1] {
+            Operand::Ident(ident) => generated::enum_from_str("Decoration", &ident)?,
+            Operand::Literal(Lit::Int(i)) => *i as u32,
+            _ => bail!("OpDecorate expected decoration"),
+        };
+        let mut parameters = Vec::new();
+        for (i, parameter) in instr.operands[2..].iter().enumerate() {
+            let parameter = match parameter {
+                Operand::IdRef(IdRef::Id(id)) => *id,
+                Operand::Literal(Lit::Int(i)) => *i as u32,
+                Operand::Ident(ident) => {
+                    let ety = generated::decorate_parameter_enum_type(decoration, i)?;
+                    let e = generated::enum_from_str(ety, &ident)?;
+                    e
+                }
+                _ => bail!("OpDecorate expected parameter"),
+            };
+            parameters.push(parameter);
+        }
+
+        let instr = InstructionBuilder::new(Op::Decorate)
+            .push(target_id)
+            .push(decoration)
+            .push_list(&parameters[..])
+            .build();
+        Ok(instr.into_words())
+    }
+    fn assemble_op_member_decorate(&mut self, instr: &Instruction) -> Result<Vec<u32>> {
+        if instr.operands.len() < 3 {
+            bail!("OpMemberDecorate expected at least 3 operands");
+        }
+        let target_id = match &instr.operands[0] {
+            Operand::IdRef(IdRef::Id(id)) => *id,
+            _ => bail!("OpMemberDecorate expected target id"),
+        };
+        let member = match &instr.operands[1] {
+            Operand::Literal(Lit::Int(i)) => *i as u32,
+            _ => bail!("OpMemberDecorate expected member"),
+        };
+        let decoration = match &instr.operands[2] {
+            Operand::Ident(ident) => generated::enum_from_str("Decoration", &ident)?,
+            Operand::Literal(Lit::Int(i)) => *i as u32,
+            _ => bail!("OpMemberDecorate expected decoration"),
+        };
+        let mut parameters = Vec::new();
+        for (i, parameter) in instr.operands[3..].iter().enumerate() {
+            let parameter = match parameter {
+                Operand::IdRef(IdRef::Id(id)) => *id,
+                Operand::Literal(Lit::Int(i)) => *i as u32,
+                Operand::Ident(ident) => {
+                    let ety = generated::decorate_parameter_enum_type(decoration, i)?;
+                    let e = generated::enum_from_str(ety, &ident)?;
+                    e
+                }
+                _ => bail!("OpMemberDecorate expected parameter"),
+            };
+            parameters.push(parameter);
+        }
+
+        let instr = InstructionBuilder::new(Op::MemberDecorate)
+            .push(target_id)
+            .push(member)
+            .push(decoration)
+            .push_list(&parameters[..])
+            .build();
+        Ok(instr.into_words())
+    }
 
     // Call this after you sanitized named refs to ID refs.
     fn assemble_special_instr(&mut self, instr: &Instruction) -> Result<Option<Vec<u32>>> {
         const OP_TYPE_INT: u32 = Op::TypeInt as u32;
         const OP_TYPE_FLOAT: u32 = Op::TypeFloat as u32;
         const OP_CONSTANT: u32 = Op::Constant as u32;
+        const OP_DECORATE: u32 = Op::Decorate as u32;
+        const OP_MEMBER_DECORATE: u32 = Op::MemberDecorate as u32;
 
         let out = match instr.opcode {
             OP_TYPE_INT => Some(self.assemble_op_type_int(instr)?),
             OP_TYPE_FLOAT => Some(self.assemble_op_type_float(instr)?),
             OP_CONSTANT => Some(self.assemble_op_constant(instr)?),
+            OP_DECORATE => Some(self.assemble_op_decorate(instr)?),
+            OP_MEMBER_DECORATE => Some(self.assemble_op_member_decorate(instr)?),
             _ => None,
         };
         Ok(out)

spirq-spvasm/src/generated/decorate_parameter_enum_type.rs

@@ -0,0 +1,78 @@
+use anyhow::{bail, Result};
+use num_traits::FromPrimitive;
+use spirq_core::spirv::Op;
+
+fn unknown_decorate_parameter_index(decoration: u32, i: usize) -> Result<&'static str> {
+    let opname = Op::from_u32(decoration)
+        .map(|op| format!("{:?}", op))
+        .unwrap_or("<unknown>".to_owned());
+    bail!(
+        "Unknown op {} ({}) parameter index: {}",
+        opname,
+        decoration,
+        i
+    )
+}
+
+pub fn decorate_parameter_enum_type(decoration: u32, i: usize) -> Result<&'static str> {
+    let out: &'static str = match decoration {
+        11 => match i {
+            0 => "BuiltIn",
+            _ => return unknown_decorate_parameter_index(decoration, i),
+        },
+        38 => match i {
+            0 => "FunctionParameterAttribute",
+            _ => return unknown_decorate_parameter_index(decoration, i),
+        },
+        39 => match i {
+            0 => "FPRoundingMode",
+            _ => return unknown_decorate_parameter_index(decoration, i),
+        },
+        40 => match i {
+            0 => "FPFastMathMode",
+            _ => return unknown_decorate_parameter_index(decoration, i),
+        },
+        41 => match i {
+            1 => "LinkageType",
+            _ => return unknown_decorate_parameter_index(decoration, i),
+        },
+        5822 => match i {
+            1 => "FPRoundingMode",
+            _ => return unknown_decorate_parameter_index(decoration, i),
+        },
+        5823 => match i {
+            1 => "FPDenormMode",
+            _ => return unknown_decorate_parameter_index(decoration, i),
+        },
+        6080 => match i {
+            1 => "FPOperationMode",
+            _ => return unknown_decorate_parameter_index(decoration, i),
+        },
+        6180 => match i {
+            0 => "AccessQualifier",
+            _ => return unknown_decorate_parameter_index(decoration, i),
+        },
+        6188 => match i {
+            0 => "HostAccessQualifier",
+            _ => return unknown_decorate_parameter_index(decoration, i),
+        },
+        6190 => match i {
+            0 => "InitializationModeQualifier",
+            _ => return unknown_decorate_parameter_index(decoration, i),
+        },
+        6442 => match i {
+            1 => "LoadCacheControl",
+            _ => return unknown_decorate_parameter_index(decoration, i),
+        },
+        6443 => match i {
+            1 => "StoreCacheControl",
+            _ => return unknown_decorate_parameter_index(decoration, i),
+        },
+        _ => bail!(
+            "{}-th parameter of decoration {} is not a enum",
+            i,
+            decoration
+        ),
+    };
+    Ok(out)
+}

spirq-spvasm/src/generated/mod.rs

@@ -1,3 +1,4 @@
+mod decorate_parameter_enum_type;
 mod enum_from_str;
 mod enum_to_str;
 mod op_from_str;
@@ -7,6 +8,7 @@ mod op_to_str;
 mod operand_enum_type;
 mod print_operand;
 
+pub use decorate_parameter_enum_type::decorate_parameter_enum_type;
 pub use enum_from_str::enum_from_str;
 pub use enum_to_str::enum_to_str;
 pub use op_from_str::op_from_str;