mirror of
https://github.com/vosen/ZLUDA.git
synced 2025-04-20 00:14:45 +00:00
Port sreg fix pass
This commit is contained in:
Andrzej Janik
parent
4e6dc07a52
commit
107f1eb17f
2 changed files with 428 additions and 2 deletions
183
ptx/src/pass/fix_special_registers.rs
Normal file
183
ptx/src/pass/fix_special_registers.rs
Normal file
|
|
@ -0,0 +1,183 @@
|
|||
use super::*;
|
||||
use std::collections::HashMap;
|
||||
|
||||
fn run<'a, 'b, 'input>(
|
||||
ptx_impl_imports: &'a mut HashMap<String, Directive<'input>>,
|
||||
typed_statements: Vec<TypedStatement>,
|
||||
numeric_id_defs: &'a mut NumericIdResolver<'b>,
|
||||
) -> Result<Vec<TypedStatement>, TranslateError> {
|
||||
let result = Vec::with_capacity(typed_statements.len());
|
||||
let mut sreg_sresolver = SpecialRegisterResolver {
|
||||
ptx_impl_imports,
|
||||
numeric_id_defs,
|
||||
result,
|
||||
};
|
||||
for statement in typed_statements {
|
||||
let statement = statement.visit_map(&mut sreg_sresolver)?;
|
||||
sreg_sresolver.result.push(statement);
|
||||
}
|
||||
Ok(sreg_sresolver.result)
|
||||
}
|
||||
|
||||
struct SpecialRegisterResolver<'a, 'b, 'input> {
|
||||
ptx_impl_imports: &'a mut HashMap<String, Directive<'input>>,
|
||||
numeric_id_defs: &'a mut NumericIdResolver<'b>,
|
||||
result: Vec<TypedStatement>,
|
||||
}
|
||||
|
||||
impl<'a, 'b, 'input> ast::VisitorMap<TypedOperand, TypedOperand, TranslateError>
|
||||
for SpecialRegisterResolver<'a, 'b, 'input>
|
||||
{
|
||||
fn visit(
|
||||
&mut self,
|
||||
operand: TypedOperand,
|
||||
_type_space: Option<(&ptx_parser::Type, ptx_parser::StateSpace)>,
|
||||
is_dst: bool,
|
||||
_relaxed_type_check: bool,
|
||||
) -> Result<TypedOperand, TranslateError> {
|
||||
operand.map(|name, vector_index| self.replace_sreg(name, is_dst, vector_index))
|
||||
}
|
||||
|
||||
fn visit_ident(
|
||||
&mut self,
|
||||
args: SpirvWord,
|
||||
_type_space: Option<(&ptx_parser::Type, ptx_parser::StateSpace)>,
|
||||
is_dst: bool,
|
||||
_relaxed_type_check: bool,
|
||||
) -> Result<SpirvWord, TranslateError> {
|
||||
self.replace_sreg(args, is_dst, None)
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, 'b, 'input> SpecialRegisterResolver<'a, 'b, 'input> {
|
||||
fn replace_sreg(
|
||||
&mut self,
|
||||
name: SpirvWord,
|
||||
is_dst: bool,
|
||||
vector_index: Option<u8>,
|
||||
) -> Result<SpirvWord, TranslateError> {
|
||||
if let Some(sreg) = self.numeric_id_defs.special_registers.get(name) {
|
||||
if is_dst {
|
||||
return Err(TranslateError::MismatchedType);
|
||||
}
|
||||
let input_arguments = match (vector_index, sreg.get_function_input_type()) {
|
||||
(Some(idx), Some(inp_type)) => {
|
||||
if inp_type != ast::ScalarType::U8 {
|
||||
return Err(TranslateError::Unreachable);
|
||||
}
|
||||
let constant = self.numeric_id_defs.register_intermediate(Some((
|
||||
ast::Type::Scalar(inp_type),
|
||||
ast::StateSpace::Reg,
|
||||
)));
|
||||
self.result.push(Statement::Constant(ConstantDefinition {
|
||||
dst: constant,
|
||||
typ: inp_type,
|
||||
value: ast::ImmediateValue::U64(idx as u64),
|
||||
}));
|
||||
vec![(
|
||||
TypedOperand::Reg(constant),
|
||||
ast::Type::Scalar(inp_type),
|
||||
ast::StateSpace::Reg,
|
||||
)]
|
||||
}
|
||||
(None, None) => Vec::new(),
|
||||
_ => return Err(TranslateError::MismatchedType),
|
||||
};
|
||||
let ocl_fn_name = [ZLUDA_PTX_PREFIX, sreg.get_unprefixed_function_name()].concat();
|
||||
let return_type = sreg.get_function_return_type();
|
||||
let fn_result = self.numeric_id_defs.register_intermediate(Some((
|
||||
ast::Type::Scalar(return_type),
|
||||
ast::StateSpace::Reg,
|
||||
)));
|
||||
let return_arguments = vec![(
|
||||
fn_result,
|
||||
ast::Type::Scalar(return_type),
|
||||
ast::StateSpace::Reg,
|
||||
)];
|
||||
let fn_call = register_external_fn_call(
|
||||
self.numeric_id_defs,
|
||||
self.ptx_impl_imports,
|
||||
ocl_fn_name.to_string(),
|
||||
return_arguments.iter().map(|(_, typ, space)| (typ, *space)),
|
||||
input_arguments.iter().map(|(_, typ, space)| (typ, *space)),
|
||||
)?;
|
||||
let data = ast::CallDetails {
|
||||
uniform: false,
|
||||
return_arguments: return_arguments
|
||||
.iter()
|
||||
.map(|(_, typ, space)| (typ.clone(), *space))
|
||||
.collect(),
|
||||
input_arguments: input_arguments
|
||||
.iter()
|
||||
.map(|(_, typ, space)| (typ.clone(), *space))
|
||||
.collect(),
|
||||
};
|
||||
let arguments = ast::CallArgs {
|
||||
return_arguments: return_arguments.iter().map(|(name, _, _)| *name).collect(),
|
||||
func: fn_call,
|
||||
input_arguments: input_arguments.iter().map(|(name, _, _)| *name).collect(),
|
||||
};
|
||||
self.result
|
||||
.push(Statement::Instruction(ast::Instruction::Call {
|
||||
data,
|
||||
arguments,
|
||||
}));
|
||||
Ok(fn_result)
|
||||
} else {
|
||||
Ok(name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn register_external_fn_call<'a>(
|
||||
id_defs: &mut NumericIdResolver,
|
||||
ptx_impl_imports: &mut HashMap<String, Directive>,
|
||||
name: String,
|
||||
return_arguments: impl Iterator<Item = (&'a ast::Type, ast::StateSpace)>,
|
||||
input_arguments: impl Iterator<Item = (&'a ast::Type, ast::StateSpace)>,
|
||||
) -> Result<SpirvWord, TranslateError> {
|
||||
match ptx_impl_imports.entry(name) {
|
||||
hash_map::Entry::Vacant(entry) => {
|
||||
let fn_id = id_defs.register_intermediate(None);
|
||||
let return_arguments = fn_arguments_to_variables(id_defs, return_arguments);
|
||||
let input_arguments = fn_arguments_to_variables(id_defs, input_arguments);
|
||||
let func_decl = ast::MethodDeclaration::<SpirvWord> {
|
||||
return_arguments,
|
||||
name: ast::MethodName::Func(fn_id),
|
||||
input_arguments,
|
||||
shared_mem: None,
|
||||
};
|
||||
let func = Function {
|
||||
func_decl: Rc::new(RefCell::new(func_decl)),
|
||||
globals: Vec::new(),
|
||||
body: None,
|
||||
import_as: Some(entry.key().clone()),
|
||||
tuning: Vec::new(),
|
||||
linkage: ast::LinkingDirective::EXTERN,
|
||||
};
|
||||
entry.insert(Directive::Method(func));
|
||||
Ok(fn_id)
|
||||
}
|
||||
hash_map::Entry::Occupied(entry) => match entry.get() {
|
||||
Directive::Method(Function { func_decl, .. }) => match (**func_decl).borrow().name {
|
||||
ast::MethodName::Func(fn_id) => Ok(fn_id),
|
||||
ast::MethodName::Kernel(_) => Err(error_unreachable()),
|
||||
},
|
||||
_ => Err(error_unreachable()),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn fn_arguments_to_variables<'a>(
|
||||
id_defs: &mut NumericIdResolver,
|
||||
args: impl Iterator<Item = (&'a ast::Type, ast::StateSpace)>,
|
||||
) -> Vec<ast::Variable<SpirvWord>> {
|
||||
args.map(|(typ, space)| ast::Variable {
|
||||
align: None,
|
||||
v_type: typ.clone(),
|
||||
state_space: space,
|
||||
name: id_defs.register_intermediate(None),
|
||||
array_init: Vec::new(),
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
}
|
||||
|
|
@ -9,6 +9,7 @@ use std::{
|
|||
};
|
||||
|
||||
mod convert_to_typed;
|
||||
mod fix_special_registers;
|
||||
mod normalize_identifiers;
|
||||
mod normalize_predicates;
|
||||
|
||||
|
|
@ -735,6 +736,235 @@ enum Statement<I, P: ast::Operand> {
|
|||
FunctionPointer(FunctionPointerDetails),
|
||||
}
|
||||
|
||||
impl<T: ast::Operand<Ident = SpirvWord>> Statement<ast::Instruction<T>, T> {
|
||||
fn visit_map<To: ast::Operand<Ident = SpirvWord>, Err>(
|
||||
self,
|
||||
visitor: &mut impl ast::VisitorMap<T, To, Err>,
|
||||
) -> std::result::Result<Statement<ast::Instruction<To>, T>, Err> {
|
||||
Ok(match self {
|
||||
Statement::Instruction(i) => {
|
||||
return ast::visit_map(i, visitor).map(Statement::Instruction)
|
||||
}
|
||||
Statement::Label(label) => {
|
||||
Statement::Label(visitor.visit_ident(label, None, false, false)?)
|
||||
}
|
||||
Statement::Variable(var) => {
|
||||
let name = visitor.visit_ident(
|
||||
var.name,
|
||||
Some((&var.v_type, var.state_space)),
|
||||
true,
|
||||
false,
|
||||
)?;
|
||||
Statement::Variable(ast::Variable {
|
||||
align: var.align,
|
||||
v_type: var.v_type,
|
||||
state_space: var.state_space,
|
||||
name,
|
||||
array_init: var.array_init,
|
||||
})
|
||||
}
|
||||
Statement::Conditional(conditional) => {
|
||||
let predicate = visitor.visit_ident(conditional.predicate, None, false, false)?;
|
||||
let if_true = visitor.visit_ident(conditional.if_true, None, false, false)?;
|
||||
let if_false = visitor.visit_ident(conditional.if_false, None, false, false)?;
|
||||
Statement::Conditional(BrachCondition {
|
||||
predicate,
|
||||
if_true,
|
||||
if_false,
|
||||
})
|
||||
}
|
||||
Statement::LoadVar(LoadVarDetails {
|
||||
arg,
|
||||
typ,
|
||||
member_index,
|
||||
}) => {
|
||||
let dst = visitor.visit_ident(
|
||||
arg.dst,
|
||||
Some((&typ, ast::StateSpace::Reg)),
|
||||
true,
|
||||
false,
|
||||
)?;
|
||||
let src = visitor.visit_ident(
|
||||
arg.src,
|
||||
Some((&typ, ast::StateSpace::Local)),
|
||||
false,
|
||||
false,
|
||||
)?;
|
||||
Statement::LoadVar(LoadVarDetails {
|
||||
arg: ast::LdArgs { dst, src },
|
||||
typ,
|
||||
member_index,
|
||||
})
|
||||
}
|
||||
Statement::StoreVar(StoreVarDetails {
|
||||
arg,
|
||||
typ,
|
||||
member_index,
|
||||
}) => {
|
||||
let src1 = visitor.visit_ident(
|
||||
arg.src1,
|
||||
Some((&typ, ast::StateSpace::Local)),
|
||||
false,
|
||||
false,
|
||||
)?;
|
||||
let src2 = visitor.visit_ident(
|
||||
arg.src2,
|
||||
Some((&typ, ast::StateSpace::Reg)),
|
||||
false,
|
||||
false,
|
||||
)?;
|
||||
Statement::StoreVar(StoreVarDetails {
|
||||
arg: ast::StArgs { src1, src2 },
|
||||
typ,
|
||||
member_index,
|
||||
})
|
||||
}
|
||||
Statement::Conversion(ImplicitConversion {
|
||||
src,
|
||||
dst,
|
||||
from_type,
|
||||
to_type,
|
||||
from_space,
|
||||
to_space,
|
||||
kind,
|
||||
}) => {
|
||||
let dst = visitor.visit_ident(
|
||||
dst,
|
||||
Some((&to_type, ast::StateSpace::Reg)),
|
||||
true,
|
||||
false,
|
||||
)?;
|
||||
let src = visitor.visit_ident(
|
||||
src,
|
||||
Some((&from_type, ast::StateSpace::Reg)),
|
||||
false,
|
||||
false,
|
||||
)?;
|
||||
Statement::Conversion(ImplicitConversion {
|
||||
src,
|
||||
dst,
|
||||
from_type,
|
||||
to_type,
|
||||
from_space,
|
||||
to_space,
|
||||
kind,
|
||||
})
|
||||
}
|
||||
Statement::Constant(ConstantDefinition { dst, typ, value }) => {
|
||||
let dst = visitor.visit_ident(
|
||||
dst,
|
||||
Some((&typ.into(), ast::StateSpace::Reg)),
|
||||
true,
|
||||
false,
|
||||
)?;
|
||||
Statement::Constant(ConstantDefinition { dst, typ, value })
|
||||
}
|
||||
Statement::RetValue(data, value) => {
|
||||
// TODO:
|
||||
// We should report type here
|
||||
let value = visitor.visit_ident(value, None, false, false)?;
|
||||
Statement::RetValue(data, value)
|
||||
}
|
||||
Statement::PtrAccess(PtrAccess {
|
||||
underlying_type,
|
||||
state_space,
|
||||
dst,
|
||||
ptr_src,
|
||||
offset_src,
|
||||
}) => {
|
||||
let dst =
|
||||
visitor.visit_ident(dst, Some((&underlying_type, state_space)), true, false)?;
|
||||
let ptr_src = visitor.visit_ident(
|
||||
ptr_src,
|
||||
Some((&underlying_type, state_space)),
|
||||
false,
|
||||
false,
|
||||
)?;
|
||||
Statement::PtrAccess(PtrAccess {
|
||||
underlying_type,
|
||||
state_space,
|
||||
dst,
|
||||
ptr_src,
|
||||
offset_src,
|
||||
})
|
||||
}
|
||||
Statement::RepackVector(RepackVectorDetails {
|
||||
is_extract,
|
||||
typ,
|
||||
packed,
|
||||
unpacked,
|
||||
relaxed_type_check,
|
||||
}) => {
|
||||
let (packed, unpacked) = if is_extract {
|
||||
let unpacked = unpacked
|
||||
.into_iter()
|
||||
.map(|ident| {
|
||||
visitor.visit_ident(
|
||||
ident,
|
||||
Some((&typ.into(), ast::StateSpace::Reg)),
|
||||
true,
|
||||
relaxed_type_check,
|
||||
)
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
let packed = visitor.visit_ident(
|
||||
packed,
|
||||
Some((
|
||||
&ast::Type::Vector(typ, unpacked.len() as u8),
|
||||
ast::StateSpace::Reg,
|
||||
)),
|
||||
false,
|
||||
false,
|
||||
)?;
|
||||
(packed, unpacked)
|
||||
} else {
|
||||
let packed = visitor.visit_ident(
|
||||
packed,
|
||||
Some((
|
||||
&ast::Type::Vector(typ, unpacked.len() as u8),
|
||||
ast::StateSpace::Reg,
|
||||
)),
|
||||
true,
|
||||
false,
|
||||
)?;
|
||||
let unpacked = unpacked
|
||||
.into_iter()
|
||||
.map(|ident| {
|
||||
visitor.visit_ident(
|
||||
ident,
|
||||
Some((&typ.into(), ast::StateSpace::Reg)),
|
||||
false,
|
||||
relaxed_type_check,
|
||||
)
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
(packed, unpacked)
|
||||
};
|
||||
Statement::RepackVector(RepackVectorDetails {
|
||||
is_extract,
|
||||
typ,
|
||||
packed,
|
||||
unpacked,
|
||||
relaxed_type_check,
|
||||
})
|
||||
}
|
||||
Statement::FunctionPointer(FunctionPointerDetails { dst, src }) => {
|
||||
let dst = visitor.visit_ident(
|
||||
dst,
|
||||
Some((
|
||||
&ast::Type::Scalar(ast::ScalarType::U64),
|
||||
ast::StateSpace::Reg,
|
||||
)),
|
||||
true,
|
||||
false,
|
||||
)?;
|
||||
let src = visitor.visit_ident(src, None, false, false)?;
|
||||
Statement::FunctionPointer(FunctionPointerDetails { dst, src })
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct BrachCondition {
|
||||
predicate: SpirvWord,
|
||||
if_true: SpirvWord,
|
||||
|
|
@ -743,7 +973,6 @@ struct BrachCondition {
|
|||
struct LoadVarDetails {
|
||||
arg: ast::LdArgs<SpirvWord>,
|
||||
typ: ast::Type,
|
||||
state_space: ast::StateSpace,
|
||||
// (index, vector_width)
|
||||
// HACK ALERT
|
||||
// For some reason IGC explodes when you try to load from builtin vectors
|
||||
|
|
@ -798,7 +1027,7 @@ struct RepackVectorDetails {
|
|||
typ: ast::ScalarType,
|
||||
packed: SpirvWord,
|
||||
unpacked: Vec<SpirvWord>,
|
||||
relaxed_type_check: bool
|
||||
relaxed_type_check: bool,
|
||||
}
|
||||
|
||||
struct FunctionPointerDetails {
|
||||
|
|
@ -876,6 +1105,20 @@ enum TypedOperand {
|
|||
VecMember(SpirvWord, u8),
|
||||
}
|
||||
|
||||
impl TypedOperand {
|
||||
fn map<Err>(
|
||||
self,
|
||||
fn_: impl FnOnce(SpirvWord, Option<u8>) -> Result<SpirvWord, Err>,
|
||||
) -> Result<Self, Err> {
|
||||
Ok(match self {
|
||||
TypedOperand::Reg(reg) => TypedOperand::Reg(fn_(reg, None)?),
|
||||
TypedOperand::RegOffset(reg, off) => TypedOperand::RegOffset(fn_(reg, None)?, off),
|
||||
TypedOperand::Imm(imm) => TypedOperand::Imm(imm),
|
||||
TypedOperand::VecMember(reg, idx) => TypedOperand::VecMember(fn_(reg, Some(idx))?, idx),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl ast::Operand for TypedOperand {
|
||||
type Ident = SpirvWord;
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue