use std::{
collections::{BTreeMap, BTreeSet},
fmt, fs, io,
path::{Path, PathBuf},
};
use super::{DeclarativeRegistryError, DialectRegistry, OperationShape};
#[derive(Clone, Debug, serde::Deserialize)]
#[serde(deny_unknown_fields)]
pub struct RegistryConfig {
#[serde(default)]
pub presets: Vec<String>,
pub builtins: Vec<String>,
pub operation_shapes: Vec<OperationShapeConfig>,
#[serde(default)]
pub operation_formats: Vec<OperationFormatConfig>,
#[serde(default)]
pub operation_alternatives: Vec<OperationAlternativesConfig>,
}
#[derive(Clone, Debug, serde::Deserialize)]
#[serde(deny_unknown_fields)]
pub struct OperationShapeConfig {
pub name: String,
pub shape: OperationShape,
}
#[derive(Clone, Debug, serde::Deserialize)]
#[serde(deny_unknown_fields)]
pub struct OperationFormatConfig {
pub name: String,
pub format: String,
}
#[derive(Clone, Debug, serde::Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct OperationAlternativesConfig {
pub name: String,
pub alternatives: Vec<OperationGrammarConfig>,
}
#[derive(Clone, Debug, serde::Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct OperationGrammarConfig {
#[serde(default)]
pub shape: Option<OperationShape>,
#[serde(default)]
pub format: Option<String>,
}
type ExpandedRegistry = (
BTreeSet<String>,
BTreeMap<String, OperationShape>,
BTreeMap<String, String>,
BTreeMap<String, Vec<OperationGrammarConfig>>,
);
const PRESET_NAMES: &[&str] = &[
"stablehlo",
"tosa",
"scf",
"linalg",
"acc",
"affine",
"amdgpu",
"amx",
"arith",
"arm_neon",
"arm_sme",
"arm_sve",
"async",
"bufferization",
"cf",
"complex",
"dlti",
"emitc",
"func",
"gpu",
"index",
"irdl",
"llvm",
"math",
"memref",
"ml_program",
"mpi",
"shard",
"nvgpu",
"nvvm",
"omp",
"pdl",
"pdl_interp",
"ptr",
"quant",
"rocdl",
"shape",
"sparse_tensor",
"smt",
"spirv",
"tensor",
"transform",
"ub",
"vector",
"wasmssa",
"x86vector",
"xegpu",
"xevm",
];
fn preset_json(name: &str) -> Option<&'static str> {
match name {
"stablehlo" => Some(include_str!("../../registries/stablehlo.json")),
"tosa" => Some(include_str!("../../registries/tosa.json")),
"scf" => Some(include_str!("../../registries/scf.json")),
"linalg" => Some(include_str!("../../registries/linalg.json")),
"acc" => Some(include_str!("../../registries/acc.json")),
"affine" => Some(include_str!("../../registries/affine.json")),
"amdgpu" => Some(include_str!("../../registries/amdgpu.json")),
"amx" => Some(include_str!("../../registries/amx.json")),
"arith" => Some(include_str!("../../registries/arith.json")),
"arm_neon" => Some(include_str!("../../registries/arm_neon.json")),
"arm_sme" => Some(include_str!("../../registries/arm_sme.json")),
"arm_sve" => Some(include_str!("../../registries/arm_sve.json")),
"async" => Some(include_str!("../../registries/async.json")),
"bufferization" => Some(include_str!("../../registries/bufferization.json")),
"cf" => Some(include_str!("../../registries/cf.json")),
"complex" => Some(include_str!("../../registries/complex.json")),
"dlti" => Some(include_str!("../../registries/dlti.json")),
"emitc" => Some(include_str!("../../registries/emitc.json")),
"func" => Some(include_str!("../../registries/func.json")),
"gpu" => Some(include_str!("../../registries/gpu.json")),
"index" => Some(include_str!("../../registries/index.json")),
"irdl" => Some(include_str!("../../registries/irdl.json")),
"llvm" => Some(include_str!("../../registries/llvm.json")),
"math" => Some(include_str!("../../registries/math.json")),
"memref" => Some(include_str!("../../registries/memref.json")),
"ml_program" => Some(include_str!("../../registries/ml_program.json")),
"mpi" => Some(include_str!("../../registries/mpi.json")),
"shard" => Some(include_str!("../../registries/shard.json")),
"nvgpu" => Some(include_str!("../../registries/nvgpu.json")),
"nvvm" => Some(include_str!("../../registries/nvvm.json")),
"omp" => Some(include_str!("../../registries/omp.json")),
"pdl" => Some(include_str!("../../registries/pdl.json")),
"pdl_interp" => Some(include_str!("../../registries/pdl_interp.json")),
"ptr" => Some(include_str!("../../registries/ptr.json")),
"quant" => Some(include_str!("../../registries/quant.json")),
"rocdl" => Some(include_str!("../../registries/rocdl.json")),
"shape" => Some(include_str!("../../registries/shape.json")),
"sparse_tensor" => Some(include_str!("../../registries/sparse_tensor.json")),
"smt" => Some(include_str!("../../registries/smt.json")),
"spirv" => Some(include_str!("../../registries/spirv.json")),
"tensor" => Some(include_str!("../../registries/tensor.json")),
"transform" => Some(include_str!("../../registries/transform.json")),
"ub" => Some(include_str!("../../registries/ub.json")),
"vector" => Some(include_str!("../../registries/vector.json")),
"wasmssa" => Some(include_str!("../../registries/wasmssa.json")),
"x86vector" => Some(include_str!("../../registries/x86vector.json")),
"xegpu" => Some(include_str!("../../registries/xegpu.json")),
"xevm" => Some(include_str!("../../registries/xevm.json")),
_ => None,
}
}
impl RegistryConfig {
pub fn from_json(json: &str) -> Result<Self, serde_json::Error> {
let config = serde_json::from_str(json)?;
let value: serde_json::Value = serde_json::from_str(json)?;
if !value.is_object()
|| value["operation_shapes"]
.as_array()
.is_some_and(|entries| entries.iter().any(|entry| !entry.is_object()))
|| value["operation_formats"]
.as_array()
.is_some_and(|entries| entries.iter().any(|entry| !entry.is_object()))
|| value["operation_alternatives"]
.as_array()
.is_some_and(|entries| {
entries.iter().any(|entry| {
!entry.is_object()
|| entry["alternatives"].as_array().is_some_and(|grammars| {
grammars.iter().any(|grammar| !grammar.is_object())
})
})
})
{
return Err(<serde_json::Error as serde::de::Error>::custom(
"registry and operation records must be JSON objects",
));
}
Ok(config)
}
pub fn build_many(configs: &[Self]) -> Result<DialectRegistry, DeclarativeRegistryError> {
let mut builtins = BTreeSet::new();
let mut shapes = BTreeMap::new();
let mut formats = BTreeMap::new();
let mut alternatives = BTreeMap::new();
for config in configs {
let (config_builtins, config_shapes, config_formats, config_alternatives) =
config.expanded()?;
builtins.extend(config_builtins);
for (name, shape) in config_shapes {
if let Some(previous) = shapes.insert(name.clone(), shape)
&& previous != shape
{
return Err(DeclarativeRegistryError::ConflictingShape(name));
}
}
for (name, format) in config_formats {
if let Some(previous) = formats.insert(name.clone(), format.clone())
&& previous != format
{
return Err(DeclarativeRegistryError::ConflictingFormat(name));
}
}
for (name, grammars) in config_alternatives {
if let Some(previous) = alternatives.insert(name.clone(), grammars.clone())
&& previous != grammars
{
return Err(DeclarativeRegistryError::ConflictingAlternatives(name));
}
}
}
Self::build_entries(builtins, shapes, formats, alternatives)
}
pub fn build(&self) -> Result<DialectRegistry, DeclarativeRegistryError> {
let (builtins, shapes, formats, alternatives) = self.expanded()?;
Self::build_entries(builtins, shapes, formats, alternatives)
}
fn expanded(&self) -> Result<ExpandedRegistry, DeclarativeRegistryError> {
let mut seen_presets = BTreeSet::new();
let mut builtins = BTreeSet::new();
let mut shapes = BTreeMap::new();
let mut formats = BTreeMap::new();
let mut alternatives = BTreeMap::new();
for name in &self.presets {
if !seen_presets.insert(name.as_str()) {
return Err(DeclarativeRegistryError::DuplicatePreset(name.clone()));
}
let json = preset_json(name)
.ok_or_else(|| DeclarativeRegistryError::UnknownPreset(name.clone()))?;
let preset = Self::from_json(json).expect("bundled registry preset must be valid");
let (preset_builtins, preset_shapes, preset_formats, preset_alternatives) =
preset.expanded()?;
builtins.extend(preset_builtins);
for (name, shape) in preset_shapes {
if let Some(previous) = shapes.insert(name.clone(), shape)
&& previous != shape
{
return Err(DeclarativeRegistryError::ConflictingShape(name));
}
}
for (name, format) in preset_formats {
formats.insert(name, format);
}
alternatives.extend(preset_alternatives);
}
let mut explicit_builtins = BTreeSet::new();
for name in &self.builtins {
if !explicit_builtins.insert(name.as_str()) {
return Err(DeclarativeRegistryError::DuplicateOperation(name.clone()));
}
builtins.insert(name.clone());
}
let mut explicit_shapes = BTreeSet::new();
for operation in &self.operation_shapes {
if !explicit_shapes.insert(operation.name.as_str()) {
return Err(DeclarativeRegistryError::DuplicateOperation(
operation.name.clone(),
));
}
if let Some(previous) = shapes.insert(operation.name.clone(), operation.shape)
&& previous != operation.shape
{
return Err(DeclarativeRegistryError::ConflictingShape(
operation.name.clone(),
));
}
}
let mut explicit_formats = BTreeSet::new();
for operation in &self.operation_formats {
if !explicit_formats.insert(operation.name.as_str())
|| explicit_shapes.contains(operation.name.as_str())
{
return Err(DeclarativeRegistryError::DuplicateOperation(
operation.name.clone(),
));
}
if let Some(previous) = formats.insert(operation.name.clone(), operation.format.clone())
&& previous != operation.format
{
return Err(DeclarativeRegistryError::ConflictingFormat(
operation.name.clone(),
));
}
}
let mut explicit_alternatives = BTreeSet::new();
for operation in &self.operation_alternatives {
if !explicit_alternatives.insert(operation.name.as_str())
|| explicit_shapes.contains(operation.name.as_str())
|| explicit_formats.contains(operation.name.as_str())
{
return Err(DeclarativeRegistryError::DuplicateOperation(
operation.name.clone(),
));
}
if operation.alternatives.len() < 2
|| operation.alternatives.iter().any(|grammar| {
matches!(
(&grammar.shape, &grammar.format),
(None, None) | (Some(_), Some(_))
)
})
{
return Err(DeclarativeRegistryError::InvalidOperationAlternatives(
operation.name.clone(),
));
}
if let Some(previous) =
alternatives.insert(operation.name.clone(), operation.alternatives.clone())
&& previous != operation.alternatives
{
return Err(DeclarativeRegistryError::ConflictingAlternatives(
operation.name.clone(),
));
}
}
Ok((builtins, shapes, formats, alternatives))
}
fn build_entries(
builtins: BTreeSet<String>,
shapes: BTreeMap<String, OperationShape>,
formats: BTreeMap<String, String>,
alternatives: BTreeMap<String, Vec<OperationGrammarConfig>>,
) -> Result<DialectRegistry, DeclarativeRegistryError> {
let builtins = builtins.iter().map(String::as_str).collect::<Vec<_>>();
let shapes = shapes
.iter()
.map(|(name, shape)| (name.as_str(), *shape))
.collect::<Vec<_>>();
let formats = formats
.iter()
.map(|(name, format)| (name.as_str(), format.as_str()))
.collect::<Vec<_>>();
DialectRegistry::declarative(&builtins)?
.extend_operation_shapes(&shapes)?
.extend_operation_formats(&formats)
.and_then(|registry| {
registry
.extend_operation_alternatives(&alternatives.into_iter().collect::<Vec<_>>())
})
}
}
#[derive(Debug)]
pub enum RegistryConfigError {
Io {
path: PathBuf,
error: io::Error,
},
Json {
path: PathBuf,
error: serde_json::Error,
},
Registry(DeclarativeRegistryError),
}
impl fmt::Display for RegistryConfigError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Io { path, error } => write!(f, "{}: {error}", path.display()),
Self::Json { path, error } => write!(f, "{}: {error}", path.display()),
Self::Registry(error) => error.fmt(f),
}
}
}
impl std::error::Error for RegistryConfigError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
Some(match self {
Self::Io { error, .. } => error,
Self::Json { error, .. } => error,
Self::Registry(error) => error,
})
}
}
impl DialectRegistry {
pub const fn preset_names() -> &'static [&'static str] {
PRESET_NAMES
}
pub fn from_name(name: &str) -> Result<Self, DeclarativeRegistryError> {
RegistryConfig {
presets: vec![name.to_owned()],
builtins: Vec::new(),
operation_shapes: Vec::new(),
operation_formats: Vec::new(),
operation_alternatives: Vec::new(),
}
.build()
}
pub fn from_config_file(path: impl AsRef<Path>) -> Result<Self, RegistryConfigError> {
Self::from_config_files([path])
}
pub fn from_config_files<P: AsRef<Path>>(
paths: impl IntoIterator<Item = P>,
) -> Result<Self, RegistryConfigError> {
let mut configs = Vec::new();
for path in paths {
let path = path.as_ref();
let json = fs::read_to_string(path).map_err(|error| RegistryConfigError::Io {
path: path.to_owned(),
error,
})?;
configs.push(RegistryConfig::from_json(&json).map_err(|error| {
RegistryConfigError::Json {
path: path.to_owned(),
error,
}
})?);
}
RegistryConfig::build_many(&configs).map_err(RegistryConfigError::Registry)
}
}