Skip to content

Commit

Permalink
Add relaxed type check information to visitors
Browse files Browse the repository at this point in the history
  • Loading branch information
vosen committed Aug 24, 2024
1 parent 7ea990e commit 69175d2
Show file tree
Hide file tree
Showing 3 changed files with 76 additions and 36 deletions.
82 changes: 56 additions & 26 deletions ptx_parser/src/ast.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,10 @@ ptx_parser_macros::generate_instruction_type!(
type: { &data.typ },
data: LdDetails,
arguments<T>: {
dst: T,
dst: {
repr: T,
relaxed_type_check: true,
},
src: {
repr: T,
space: { data.state_space },
Expand All @@ -51,7 +54,10 @@ ptx_parser_macros::generate_instruction_type!(
repr: T,
space: { data.state_space },
},
src2: T,
src2: {
repr: T,
relaxed_type_check: true,
}
}
},
Mul {
Expand Down Expand Up @@ -157,10 +163,13 @@ ptx_parser_macros::generate_instruction_type!(
dst: {
repr: T,
type: { Type::Scalar(data.to) },
// TODO: double check
relaxed_type_check: true,
},
src: {
repr: T,
type: { Type::Scalar(data.from) },
relaxed_type_check: true,
},
}
},
Expand Down Expand Up @@ -494,34 +503,43 @@ pub trait Visitor<T: Operand, Err> {
args: &T,
type_space: Option<(&Type, StateSpace)>,
is_dst: bool,
relaxed_type_check: bool,
) -> Result<(), Err>;
fn visit_ident(
&mut self,
args: &T::Ident,
type_space: Option<(&Type, StateSpace)>,
is_dst: bool,
relaxed_type_check: bool,
) -> Result<(), Err>;
}

impl<T: Operand, Err, Fn: FnMut(&T, Option<(&Type, StateSpace)>, bool) -> Result<(), Err>>
impl<T: Operand, Err, Fn: FnMut(&T, Option<(&Type, StateSpace)>, bool, bool) -> Result<(), Err>>
Visitor<T, Err> for Fn
{
fn visit(
&mut self,
args: &T,
type_space: Option<(&Type, StateSpace)>,
is_dst: bool,
relaxed_type_check: bool,
) -> Result<(), Err> {
(self)(args, type_space, is_dst)
(self)(args, type_space, is_dst, relaxed_type_check)
}

fn visit_ident(
&mut self,
args: &T::Ident,
type_space: Option<(&Type, StateSpace)>,
is_dst: bool,
relaxed_type_check: bool,
) -> Result<(), Err> {
(self)(&T::from_ident(*args), type_space, is_dst)
(self)(
&T::from_ident(*args),
type_space,
is_dst,
relaxed_type_check,
)
}
}

Expand All @@ -531,12 +549,14 @@ pub trait VisitorMut<T: Operand, Err> {
args: &mut T,
type_space: Option<(&Type, StateSpace)>,
is_dst: bool,
relaxed_type_check: bool,
) -> Result<(), Err>;
fn visit_ident(
&mut self,
args: &mut T::Ident,
type_space: Option<(&Type, StateSpace)>,
is_dst: bool,
relaxed_type_check: bool,
) -> Result<(), Err>;
}

Expand All @@ -546,37 +566,44 @@ pub trait VisitorMap<From: Operand, To: Operand, Err> {
args: From,
type_space: Option<(&Type, StateSpace)>,
is_dst: bool,
relaxed_type_check: bool,
) -> Result<To, Err>;
fn visit_ident(
&mut self,
args: From::Ident,
type_space: Option<(&Type, StateSpace)>,
is_dst: bool,
relaxed_type_check: bool,
) -> Result<To::Ident, Err>;
}

impl<T: Copy, U: Copy, Err, Fn> VisitorMap<ParsedOperand<T>, ParsedOperand<U>, Err> for Fn
where
Fn: FnMut(T, Option<(&Type, StateSpace)>, bool) -> Result<U, Err>,
Fn: FnMut(T, Option<(&Type, StateSpace)>, bool, bool) -> Result<U, Err>,
{
fn visit(
&mut self,
args: ParsedOperand<T>,
type_space: Option<(&Type, StateSpace)>,
is_dst: bool,
relaxed_type_check: bool,
) -> Result<ParsedOperand<U>, Err> {
Ok(match args {
ParsedOperand::Reg(ident) => ParsedOperand::Reg((self)(ident, type_space, is_dst)?),
ParsedOperand::RegOffset(ident, imm) => {
ParsedOperand::RegOffset((self)(ident, type_space, is_dst)?, imm)
ParsedOperand::Reg(ident) => {
ParsedOperand::Reg((self)(ident, type_space, is_dst, relaxed_type_check)?)
}
ParsedOperand::RegOffset(ident, imm) => ParsedOperand::RegOffset(
(self)(ident, type_space, is_dst, relaxed_type_check)?,
imm,
),
ParsedOperand::Imm(imm) => ParsedOperand::Imm(imm),
ParsedOperand::VecMember(ident, index) => {
ParsedOperand::VecMember((self)(ident, type_space, is_dst)?, index)
}
ParsedOperand::VecMember(ident, index) => ParsedOperand::VecMember(
(self)(ident, type_space, is_dst, relaxed_type_check)?,
index,
),
ParsedOperand::VecPack(vec) => ParsedOperand::VecPack(
vec.into_iter()
.map(|ident| (self)(ident, type_space, is_dst))
.map(|ident| (self)(ident, type_space, is_dst, relaxed_type_check))
.collect::<Result<Vec<_>, _>>()?,
),
})
Expand All @@ -587,31 +614,34 @@ where
args: T,
type_space: Option<(&Type, StateSpace)>,
is_dst: bool,
relaxed_type_check: bool,
) -> Result<U, Err> {
(self)(args, type_space, is_dst)
(self)(args, type_space, is_dst, relaxed_type_check)
}
}

impl<T: Operand<Ident = T>, U: Operand<Ident = U>, Err, Fn> VisitorMap<T, U, Err> for Fn
where
Fn: FnMut(T, Option<(&Type, StateSpace)>, bool) -> Result<U, Err>,
Fn: FnMut(T, Option<(&Type, StateSpace)>, bool, bool) -> Result<U, Err>,
{
fn visit(
&mut self,
args: T,
type_space: Option<(&Type, StateSpace)>,
is_dst: bool,
relaxed_type_check: bool,
) -> Result<U, Err> {
(self)(args, type_space, is_dst)
(self)(args, type_space, is_dst, relaxed_type_check)
}

fn visit_ident(
&mut self,
args: T,
type_space: Option<(&Type, StateSpace)>,
is_dst: bool,
relaxed_type_check: bool,
) -> Result<U, Err> {
(self)(args, type_space, is_dst)
(self)(args, type_space, is_dst, relaxed_type_check)
}
}

Expand Down Expand Up @@ -1198,15 +1228,15 @@ impl<T: Operand> CallArgs<T> {
.iter()
.zip(details.return_arguments.iter())
{
visitor.visit_ident(param, Some((type_, *space)), true)?;
visitor.visit_ident(param, Some((type_, *space)), true, false)?;
}
visitor.visit_ident(&self.func, None, false)?;
visitor.visit_ident(&self.func, None, false, false)?;
for (param, (type_, space)) in self
.input_arguments
.iter()
.zip(details.input_arguments.iter())
{
visitor.visit(param, Some((type_, *space)), true)?;
visitor.visit(param, Some((type_, *space)), true, false)?;
}
Ok(())
}
Expand All @@ -1222,15 +1252,15 @@ impl<T: Operand> CallArgs<T> {
.iter_mut()
.zip(details.return_arguments.iter())
{
visitor.visit_ident(param, Some((type_, *space)), true)?;
visitor.visit_ident(param, Some((type_, *space)), true, false)?;
}
visitor.visit_ident(&mut self.func, None, false)?;
visitor.visit_ident(&mut self.func, None, false, false)?;
for (param, (type_, space)) in self
.input_arguments
.iter_mut()
.zip(details.input_arguments.iter())
{
visitor.visit(param, Some((type_, *space)), true)?;
visitor.visit(param, Some((type_, *space)), true, false)?;
}
Ok(())
}
Expand All @@ -1245,14 +1275,14 @@ impl<T: Operand> CallArgs<T> {
.return_arguments
.into_iter()
.zip(details.return_arguments.iter())
.map(|(param, (type_, space))| visitor.visit_ident(param, Some((type_, *space)), true))
.map(|(param, (type_, space))| visitor.visit_ident(param, Some((type_, *space)), true, false))
.collect::<Result<Vec<_>, _>>()?;
let func = visitor.visit_ident(self.func, None, false)?;
let func = visitor.visit_ident(self.func, None, false, false)?;
let input_arguments = self
.input_arguments
.into_iter()
.zip(details.input_arguments.iter())
.map(|(param, (type_, space))| visitor.visit(param, Some((type_, *space)), true))
.map(|(param, (type_, space))| visitor.visit(param, Some((type_, *space)), true, false))
.collect::<Result<Vec<_>, _>>()?;
Ok(CallArgs {
return_arguments,
Expand Down
2 changes: 1 addition & 1 deletion ptx_parser_macros/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1017,7 +1017,7 @@ pub fn generate_instruction_type(tokens: proc_macro::TokenStream) -> proc_macro:
input.emit_arg_types(&mut result);
input.emit_instruction_type(&mut result);
input.emit_visit(&mut result);
//input.emit_visit_mut(&mut result);
input.emit_visit_mut(&mut result);
input.emit_visit_map(&mut result);
result.into()
}
28 changes: 19 additions & 9 deletions ptx_parser_macros_impl/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -512,12 +512,13 @@ pub struct ArgumentField {
pub repr: Type,
pub space: Option<Expr>,
pub type_: Option<Expr>,
pub relaxed_type_check: bool,
}

impl ArgumentField {
fn parse_block(
input: syn::parse::ParseStream,
) -> syn::Result<(Type, Option<Expr>, Option<Expr>, Option<bool>)> {
) -> syn::Result<(Type, Option<Expr>, Option<Expr>, Option<bool>, bool)> {
let content;
braced!(content in input);
let all_fields =
Expand All @@ -531,6 +532,9 @@ impl ArgumentField {
let name_ident = content.parse::<Ident>()?;
content.parse::<Token![:]>()?;
match &*name_ident.to_string() {
"relaxed_type_check" => {
ExprOrPath::RelaxedTypeCheck(content.parse::<LitBool>()?.value)
}
"repr" => ExprOrPath::Repr(content.parse::<Type>()?),
"space" => ExprOrPath::Space(content.parse::<Expr>()?),
"dst" => {
Expand All @@ -552,15 +556,17 @@ impl ArgumentField {
let mut type_ = None;
let mut space = None;
let mut is_dst = None;
let mut relaxed_type_check = false;
for exp_or_path in all_fields {
match exp_or_path {
ExprOrPath::Repr(r) => repr = Some(r),
ExprOrPath::Type(t) => type_ = Some(t),
ExprOrPath::Space(s) => space = Some(s),
ExprOrPath::Dst(x) => is_dst = Some(x),
ExprOrPath::RelaxedTypeCheck(relaxed) => relaxed_type_check = relaxed,
}
}
Ok((repr.unwrap(), type_, space, is_dst))
Ok((repr.unwrap(), type_, space, is_dst, relaxed_type_check))
}

fn parse_basic(input: &syn::parse::ParseBuffer) -> syn::Result<Type> {
Expand Down Expand Up @@ -605,6 +611,7 @@ impl ArgumentField {
.map(|space| quote! { #space })
.unwrap_or_else(|| quote! { StateSpace::Reg });
let is_dst = self.is_dst;
let relaxed_type_check = self.relaxed_type_check;
let name = &self.name;
let type_space = if is_typeless {
quote! {
Expand All @@ -622,14 +629,14 @@ impl ArgumentField {
quote! {
{
#type_space
visitor.visit_ident(&mut arguments.#name, type_space, #is_dst)?;
visitor.visit_ident(&mut arguments.#name, type_space, #is_dst, #relaxed_type_check)?;
}
}
} else {
quote! {
{
#type_space
visitor.visit_ident(& arguments.#name, type_space, #is_dst)?;
visitor.visit_ident(& arguments.#name, type_space, #is_dst, #relaxed_type_check)?;
}
}
}
Expand All @@ -655,7 +662,7 @@ impl ArgumentField {
};
quote! {{
#type_space
#operand_fn(#arguments_name, |x| visitor.visit(x, type_space, #is_dst))?;
#operand_fn(#arguments_name, |x| visitor.visit(x, type_space, #is_dst, #relaxed_type_check))?;
}}
}
}
Expand All @@ -679,6 +686,7 @@ impl ArgumentField {
.map(|space| quote! { #space })
.unwrap_or_else(|| quote! { StateSpace::Reg });
let is_dst = self.is_dst;
let relaxed_type_check = self.relaxed_type_check;
let name = &self.name;
let type_space = if is_typeless {
quote! {
Expand All @@ -693,11 +701,11 @@ impl ArgumentField {
};
let map_call = if is_ident {
quote! {
visitor.visit_ident(arguments.#name, type_space, #is_dst)?
visitor.visit_ident(arguments.#name, type_space, #is_dst, #relaxed_type_check)?
}
} else {
quote! {
MapOperand::map(arguments.#name, |x| visitor.visit(x, type_space, #is_dst))?
MapOperand::map(arguments.#name, |x| visitor.visit(x, type_space, #is_dst, #relaxed_type_check))?
}
};
quote! {
Expand Down Expand Up @@ -739,10 +747,10 @@ impl Parse for ArgumentField {

input.parse::<Token![:]>()?;
let lookahead = input.lookahead1();
let (repr, type_, space, is_dst) = if lookahead.peek(token::Brace) {
let (repr, type_, space, is_dst, relaxed_type_check) = if lookahead.peek(token::Brace) {
Self::parse_block(input)?
} else if lookahead.peek(syn::Ident) {
(Self::parse_basic(input)?, None, None, None)
(Self::parse_basic(input)?, None, None, None, false)
} else {
return Err(lookahead.error());
};
Expand All @@ -756,6 +764,7 @@ impl Parse for ArgumentField {
repr,
type_,
space,
relaxed_type_check
})
}
}
Expand All @@ -765,6 +774,7 @@ enum ExprOrPath {
Type(Expr),
Space(Expr),
Dst(bool),
RelaxedTypeCheck(bool),
}

#[cfg(test)]
Expand Down

0 comments on commit 69175d2

Please sign in to comment.