From f67727c3bee2e4cb8c9c6754ded2fc9f0cce49e7 Mon Sep 17 00:00:00 2001 From: Rik Huijzer Date: Fri, 27 Sep 2024 08:59:33 +0200 Subject: [PATCH] Refactor --- src/dialect/arith/op.rs | 37 ++++++++++++++++++++++--------------- src/dialect/func/op.rs | 24 ++++++++++++++---------- src/ir/block.rs | 2 +- src/ir/mod.rs | 2 ++ src/ir/operation.rs | 22 ++++++++++++++-------- src/ir/value.rs | 7 +++++++ 6 files changed, 60 insertions(+), 34 deletions(-) diff --git a/src/dialect/arith/op.rs b/src/dialect/arith/op.rs index ba479fa1..21f44452 100644 --- a/src/dialect/arith/op.rs +++ b/src/dialect/arith/op.rs @@ -6,6 +6,7 @@ use crate::ir::Op; use crate::ir::OpOperand; use crate::ir::OpResult; use crate::ir::OperationName; +use crate::ir::OperationOperands; use crate::ir::Type; use crate::ir::Value; use crate::parser::Parse; @@ -59,19 +60,22 @@ impl Parse for ConstantOp { let operation_name = parser.expect(TokenKind::BareIdentifier)?; assert!(operation_name.lexeme == "arith.constant"); operation.set_name(ConstantOp::operation_name()); - let operand = match parser.peek().kind { + let operand: OpOperand = match parser.peek().kind { TokenKind::Integer => { let integer = parser.advance(); let integer = integer.lexeme.clone(); let typ = parser.typ()?; - let argument = BlockArgument::new(integer, typ); - Value::BlockArgument(argument) + // Determine value. + todo!(); + // let operand = OpOperand::new( + // Value::BlockArgument(argument) } _ => { return Err(anyhow::anyhow!("Expected integer constant")); } }; - operation.set_operands(Arc::new(vec![Arc::new(operand)])); + let operand = Arc::new(RwLock::new(operand)); + operation.set_operands(Arc::new(RwLock::new(vec![operand]))); operation.set_parent(parent); let operation = Arc::new(RwLock::new(operation)); Ok(Arc::new(RwLock::new(ConstantOp { operation }))) @@ -86,6 +90,7 @@ impl AddiOp { /// Canonicalize `addi(addi(x, c0), c1) -> addi(x, c0 + c1)`. fn addi_add_constant(&mut self) -> CanonicalizeResult { let operands = self.operation().read().unwrap().operands(); + let operands = operands.read().unwrap(); assert!(operands.len() == 2); let lhs = &operands[0]; let rhs = &operands[1]; @@ -112,25 +117,27 @@ impl Op for AddiOp { } impl Parser { - fn operand(&mut self, parent: Arc>) -> Result> { + fn operand(&mut self, parent: Arc>) -> Result>> { let identifier = self.expect(TokenKind::PercentIdentifier)?; let name = identifier.lexeme.clone(); let block = parent.read().unwrap(); - let first_use = block.first_use(name.clone()); - let typ = Type::new("any".to_string()); - let value = Value::OpResult(OpResult::new(name, typ)); - let operand = OpOperand::new(value, name); - Ok(Arc::new(operand)) - } - pub fn operands(&mut self, parent: Arc>) -> Result> { + let assignment = block.assignment(name.clone()); + todo!(); + // let typ = Type::new("any".to_string()); + // let value = Value::OpResult(OpResult::new(name, typ)); + // let operand = OpOperand::new(value, name); + // Ok(Arc::new(operand)) + } + pub fn operands(&mut self, parent: Arc>) -> Result { let mut arguments = vec![]; while self.check(TokenKind::PercentIdentifier) { - arguments.push(self.operand(parent.clone())?); + let operand = self.operand(parent.clone())?; + arguments.push(operand); if self.check(TokenKind::Comma) { let _comma = self.advance(); } } - Ok(arguments) + Ok(Arc::new(RwLock::new(arguments))) } pub fn results(&mut self) -> Result>> { let mut results = vec![]; @@ -161,7 +168,7 @@ impl Parse for AddiOp { operation.set_name(AddiOp::operation_name()); assert!(parent.is_some()); operation.set_parent(parent.clone()); - operation.set_operands(Arc::new(parser.operands(parent.unwrap())?)); + operation.set_operands(parser.operands(parent.unwrap())?); let _colon = parser.expect(TokenKind::Colon)?; let result_type = parser.expect(TokenKind::IntType)?; let result_type = Type::new(result_type.lexeme.clone()); diff --git a/src/dialect/func/op.rs b/src/dialect/func/op.rs index b009fd65..85edb1cd 100644 --- a/src/dialect/func/op.rs +++ b/src/dialect/func/op.rs @@ -3,6 +3,7 @@ use crate::ir::BlockArgument; use crate::ir::Op; use crate::ir::OpResult; use crate::ir::Operation; +use crate::ir::OperationArguments; use crate::ir::OperationName; use crate::ir::Type; use crate::ir::Value; @@ -39,8 +40,10 @@ impl Op for FuncOp { .read() .unwrap() .operands() + .read() + .unwrap() .iter() - .map(|o| o.to_string()) + .map(|o| o.read().unwrap().to_string()) .collect::>() .join(", "); write!(f, "{}", joined)?; @@ -85,7 +88,7 @@ impl Parser { Ok(identifier.lexeme.clone()) } /// %arg0 : i64, - pub fn function_argument(&mut self) -> Result> { + pub fn function_argument(&mut self) -> Result>> { let identifier = self.expect(TokenKind::PercentIdentifier)?; let name = identifier.lexeme.clone(); let typ = if self.check(TokenKind::Colon) { @@ -96,14 +99,14 @@ impl Parser { Type::new("any".to_string()) }; let arg = BlockArgument::new(name, typ); - let operand: Value = Value::BlockArgument(arg); + let operand = Arc::new(RwLock::new(arg)); if self.check(TokenKind::Comma) { self.advance(); } - Ok(Arc::new(operand)) + Ok(operand) } /// Parse `(%arg0 : i64, %arg1 : i64)`. - pub fn function_arguments(&mut self) -> Result>> { + pub fn function_arguments(&mut self) -> Result { let _lparen = self.expect(TokenKind::LParen)?; let mut operands = vec![]; while self.check(TokenKind::PercentIdentifier) { @@ -115,7 +118,7 @@ impl Parser { } else { return Err(anyhow::anyhow!("Expected ')', got {:?}", self.peek().kind)); } - Ok(operands) + Ok(Arc::new(RwLock::new(operands))) } pub fn result_types(&mut self) -> Result> { let mut result_types = vec![]; @@ -150,7 +153,7 @@ impl Parse for FuncOp { let _operation_name = parser.advance(); let identifier = parser.identifier(TokenKind::AtIdentifier).unwrap(); let mut operation = Operation::default(); - operation.set_operands(Arc::new(parser.function_arguments()?)); + operation.set_arguments(parser.function_arguments()?); operation.set_result_types(parser.result_types()?); operation.set_region(parser.region()?); operation.set_parent(parent); @@ -186,8 +189,9 @@ impl Op for ReturnOp { let operation = operation.read().unwrap(); let operands = operation.operands().clone(); write!(f, "return")?; - for operand in &*operands { - write!(f, " {}", operand)?; + let operands = operands.read().unwrap(); + for operand in operands.iter() { + write!(f, " {}", operand.read().unwrap())?; } let result_types = operation.result_types(); assert!(!result_types.is_empty(), "Expected result types to be set"); @@ -211,7 +215,7 @@ impl Parse for ReturnOp { let mut operation = Operation::default(); assert!(parent.is_some()); operation.set_parent(parent.clone()); - operation.set_operands(Arc::new(parser.operands(parent.clone().unwrap())?)); + operation.set_operands(parser.operands(parent.clone().unwrap())?); let _colon = parser.expect(TokenKind::Colon)?; let return_type = parser.expect(TokenKind::IntType)?; let return_type = Type::new(return_type.lexeme.clone()); diff --git a/src/ir/block.rs b/src/ir/block.rs index 31aa2d45..a0520e2d 100644 --- a/src/ir/block.rs +++ b/src/ir/block.rs @@ -49,7 +49,7 @@ impl Block { for operand in operands.iter() { let other_operand = operand.clone(); let write = other_operand.write().unwrap(); - let value = operand.write().unwrap().value(); + let value = other_operand.write().unwrap().value(); let value = value.clone().read().unwrap(); match value { Value::BlockArgument(block_argument) => { diff --git a/src/ir/mod.rs b/src/ir/mod.rs index 5761c25c..e1c298ed 100644 --- a/src/ir/mod.rs +++ b/src/ir/mod.rs @@ -12,7 +12,9 @@ pub use crate::ir::block::Block; pub use crate::ir::module::ModuleOp; pub use crate::ir::op::Op; pub use crate::ir::operation::Operation; +pub use crate::ir::operation::OperationArguments; pub use crate::ir::operation::OperationName; +pub use crate::ir::operation::OperationOperands; pub use crate::ir::region::Region; pub use crate::ir::value::BlockArgument; pub use crate::ir::value::OpOperand; diff --git a/src/ir/operation.rs b/src/ir/operation.rs index ac2f098d..182a2cd7 100644 --- a/src/ir/operation.rs +++ b/src/ir/operation.rs @@ -25,6 +25,9 @@ impl OperationName { } } +pub type OperationArguments = Arc>>>>; +pub type OperationOperands = Arc>>>>; + /// Note that MLIR distinguishes between Operation and Op. /// Operation generically models all operations. /// Op is an interface for more specific operations. @@ -39,8 +42,8 @@ impl OperationName { pub struct Operation { name: OperationName, /// Used by `FuncOp` to store its arguments. - arguments: Arc>>>>, - operands: Arc>>>>, + arguments: OperationArguments, + operands: OperationOperands, attributes: Vec>, /// Results can be `Values`, so either `BlockArgument` or `OpResult`. results: Vec>, @@ -55,8 +58,8 @@ pub struct Operation { impl Operation { pub fn new( name: OperationName, - arguments: Arc>>>>, - operands: Arc>>>>, + arguments: OperationArguments, + operands: OperationOperands, attributes: Vec>, results: Vec>, result_types: Vec, @@ -77,13 +80,13 @@ impl Operation { pub fn name(&self) -> String { self.name.name.clone() } - pub fn arguments(&self) -> Arc>>>> { + pub fn arguments(&self) -> OperationArguments { self.arguments.clone() } - pub fn operands(&self) -> Arc>>>> { + pub fn operands(&self) -> OperationOperands { self.operands.clone() } - pub fn operands_mut(&mut self) -> &mut Arc>>>> { + pub fn operands_mut(&mut self) -> &mut OperationOperands { &mut self.operands } pub fn attributes(&self) -> Vec> { @@ -105,7 +108,10 @@ impl Operation { pub fn set_name(&mut self, name: OperationName) { self.name = name; } - pub fn set_operands(&mut self, operands: Arc>>>>) { + pub fn set_arguments(&mut self, arguments: OperationArguments) { + self.arguments = arguments; + } + pub fn set_operands(&mut self, operands: OperationOperands) { self.operands = operands; } pub fn set_attributes(&mut self, attributes: Vec>) { diff --git a/src/ir/value.rs b/src/ir/value.rs index fa313faa..da6ea9a5 100644 --- a/src/ir/value.rs +++ b/src/ir/value.rs @@ -70,6 +70,7 @@ impl Display for Value { } } } + pub struct OpOperand { pub value: Arc>, pub operand_name: String, @@ -89,3 +90,9 @@ impl OpOperand { &self.operand_name } } + +impl Display for OpOperand { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self) + } +}