├── .gitignore ├── Cargo.lock ├── Cargo.toml ├── macros ├── Cargo.toml └── src │ └── lib.rs └── playground ├── Cargo.toml └── src └── main.rs /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | -------------------------------------------------------------------------------- /Cargo.lock: -------------------------------------------------------------------------------- 1 | # This file is automatically @generated by Cargo. 2 | # It is not intended for manual editing. 3 | version = 3 4 | 5 | [[package]] 6 | name = "macros" 7 | version = "0.1.0" 8 | 9 | [[package]] 10 | name = "playground" 11 | version = "0.1.0" 12 | dependencies = [ 13 | "macros", 14 | ] 15 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [workspace] 2 | members = [ 3 | "playground", 4 | "macros", 5 | ] 6 | -------------------------------------------------------------------------------- /macros/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "macros" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | [lib] 7 | proc-macro = true 8 | 9 | [dependencies] 10 | -------------------------------------------------------------------------------- /macros/src/lib.rs: -------------------------------------------------------------------------------- 1 | use proc_macro::{Group, Ident, Punct, Spacing, Span, TokenStream, TokenTree}; 2 | 3 | //use crate::helpers::expect_punct; 4 | 5 | fn expect_group(it: &mut impl Iterator) -> Group { 6 | if let Some(TokenTree::Group(group)) = it.next() { 7 | group 8 | } else { 9 | panic!("Expected Group") 10 | } 11 | } 12 | 13 | fn expect_punct(it: &mut impl Iterator) -> String { 14 | if let Some(TokenTree::Punct(punct)) = it.next() { 15 | punct.to_string() 16 | } else { 17 | panic!("Expected Group") 18 | } 19 | } 20 | 21 | fn drop_until_punct(it: &mut impl Iterator, delimiter: &str) { 22 | for token in it.by_ref() { 23 | if let TokenTree::Punct(punct) = token { 24 | if delimiter.contains(&punct.to_string()) { 25 | break; 26 | } 27 | } 28 | } 29 | } 30 | 31 | struct VersionConfig { 32 | fields: &'static [&'static str], 33 | enums: &'static [&'static [&'static str]], 34 | versions: &'static [&'static [&'static str]], 35 | } 36 | 37 | static AGX_VERSIONS: VersionConfig = VersionConfig { 38 | fields: &["G", "V"], 39 | enums: &[&["G13G", "G13X", "G14G"], &["V12_3", "V13_0b4"]], 40 | versions: &[ 41 | &["G13G", "V12_3"], 42 | &["G13G", "V13_0b4"], 43 | &["G13X", "V12_3"], 44 | &["G13X", "V13_0b4"], 45 | &["G14G", "V13_0b4"], 46 | ], 47 | }; 48 | 49 | fn check_version( 50 | config: &VersionConfig, 51 | ver: &[usize], 52 | it: &mut impl Iterator, 53 | ) -> bool { 54 | let first = it.next().unwrap(); 55 | let val: bool = match &first { 56 | TokenTree::Group(group) => check_version(config, ver, &mut group.stream().into_iter()), 57 | TokenTree::Ident(ident) => { 58 | let key = config 59 | .fields 60 | .iter() 61 | .position(|&r| r == ident.to_string()) 62 | .unwrap_or_else(|| panic!("Unknown field {}", ident)); 63 | let mut operator = expect_punct(it); 64 | let mut rhs_token = it.next().unwrap(); 65 | if let TokenTree::Punct(punct) = &rhs_token { 66 | operator.extend(std::iter::once(punct.as_char())); 67 | rhs_token = it.next().unwrap(); 68 | } 69 | let rhs_name = if let TokenTree::Ident(ident) = &rhs_token { 70 | ident.to_string() 71 | } else { 72 | panic!("Unexpected token {}", ident) 73 | }; 74 | 75 | let rhs = config.enums[key] 76 | .iter() 77 | .position(|&r| r == rhs_name) 78 | .unwrap_or_else(|| panic!("Unknown field {}", ident)); 79 | let lhs = ver[key]; 80 | 81 | match operator.as_str() { 82 | "==" => lhs == rhs, 83 | "!=" => lhs != rhs, 84 | ">" => lhs > rhs, 85 | ">=" => lhs >= rhs, 86 | "<" => lhs < rhs, 87 | "<=" => lhs <= rhs, 88 | _ => panic!("Unknown operator {}", operator), 89 | } 90 | } 91 | _ => { 92 | panic!("Unknown token {}", first) 93 | } 94 | }; 95 | 96 | let boolop = it.next(); 97 | match boolop { 98 | Some(TokenTree::Punct(punct)) => { 99 | let right = expect_punct(it); 100 | if right != punct.to_string() { 101 | panic!("Unexpected op {}{}", punct, right); 102 | } 103 | match punct.as_char() { 104 | '&' => val && check_version(config, ver, it), 105 | '|' => val || check_version(config, ver, it), 106 | _ => panic!("Unexpected op {}{}", right, right), 107 | } 108 | } 109 | Some(a) => panic!("Unexpected op {}", a), 110 | None => val, 111 | } 112 | } 113 | 114 | fn filter_versions( 115 | config: &VersionConfig, 116 | tag: &str, 117 | ver: &[usize], 118 | tree: impl IntoIterator, 119 | is_struct: bool, 120 | ) -> Vec { 121 | let mut out = Vec::::new(); 122 | let mut it = tree.into_iter(); 123 | 124 | while let Some(token) = it.next() { 125 | match &token { 126 | TokenTree::Punct(punct) if punct.to_string() == "#" => { 127 | let group = expect_group(&mut it); 128 | let mut grp_it = group.stream().into_iter(); 129 | let attr = grp_it.next().unwrap(); 130 | match attr { 131 | TokenTree::Ident(ident) if ident.to_string() == "ver" => { 132 | if check_version(config, ver, &mut grp_it) { 133 | } else if is_struct { 134 | drop_until_punct(&mut it, ","); 135 | } else { 136 | let first = it.next().unwrap(); 137 | match &first { 138 | TokenTree::Group(_) => (), 139 | _ => { 140 | drop_until_punct(&mut it, ",;"); 141 | } 142 | } 143 | } 144 | } 145 | _ => { 146 | out.push(token.clone()); 147 | out.push(TokenTree::Group(group.clone())); 148 | } 149 | } 150 | } 151 | TokenTree::Punct(punct) if punct.to_string() == ":" => { 152 | let next = it.next(); 153 | match next { 154 | Some(TokenTree::Punct(punct)) if punct.to_string() == ":" => (), 155 | Some(a) => { 156 | out.push(token.clone()); 157 | out.push(a.clone()); 158 | continue; 159 | } 160 | None => { 161 | out.push(token.clone()); 162 | continue; 163 | } 164 | } 165 | 166 | let next = it.next(); 167 | match next { 168 | Some(TokenTree::Ident(idtag)) if idtag.to_string() == "ver" => { 169 | let ident = match out.pop() { 170 | Some(TokenTree::Ident(ident)) => ident, 171 | a => panic!("$ver not following ident: {:?}", a), 172 | }; 173 | let name = ident.to_string() + tag; 174 | let new_ident = Ident::new(name.as_str(), ident.span()); 175 | out.push(TokenTree::Ident(new_ident)); 176 | } 177 | Some(a) => { 178 | out.push(token.clone()); 179 | out.push(token.clone()); 180 | out.push(a.clone()); 181 | } 182 | None => { 183 | out.push(token.clone()); 184 | out.push(token.clone()); 185 | } 186 | } 187 | } 188 | TokenTree::Group(group) => { 189 | let new_body = 190 | filter_versions(config, tag, ver, &mut group.stream().into_iter(), is_struct); 191 | let mut stream = TokenStream::new(); 192 | stream.extend(new_body); 193 | let mut filtered_group = Group::new(group.delimiter(), stream); 194 | filtered_group.set_span(group.span()); 195 | out.push(TokenTree::Group(filtered_group)); 196 | } 197 | _ => { 198 | out.push(token.clone()); 199 | } 200 | } 201 | } 202 | 203 | out 204 | } 205 | 206 | #[proc_macro_attribute] 207 | pub fn versions(attr: TokenStream, item: TokenStream) -> TokenStream { 208 | let config = match attr.to_string().as_str() { 209 | "AGX" => &AGX_VERSIONS, 210 | _ => panic!("Unknown version group {}", attr), 211 | }; 212 | 213 | let mut it = item.into_iter(); 214 | let mut out = TokenStream::new(); 215 | let mut body: Vec = Vec::new(); 216 | let mut is_struct = false; 217 | 218 | while let Some(token) = it.next() { 219 | match token { 220 | TokenTree::Punct(punct) if punct.to_string() == "#" => { 221 | body.push(TokenTree::Punct(punct)); 222 | body.push(it.next().unwrap()); 223 | } 224 | TokenTree::Ident(ident) if ident.to_string() == "struct" => { 225 | body.push(TokenTree::Ident(ident)); 226 | body.push(it.next().unwrap()); 227 | // This isn't valid syntax in a struct definition, so add it for the user 228 | body.push(TokenTree::Punct(Punct::new(':', Spacing::Joint))); 229 | body.push(TokenTree::Punct(Punct::new(':', Spacing::Alone))); 230 | body.push(TokenTree::Ident(Ident::new("ver", Span::call_site()))); 231 | is_struct = true; 232 | break; 233 | } 234 | TokenTree::Ident(ident) if ident.to_string() == "impl" => { 235 | body.push(TokenTree::Ident(ident)); 236 | break; 237 | } 238 | TokenTree::Ident(ident) if ident.to_string() == "fn" => { 239 | body.push(TokenTree::Ident(ident)); 240 | break; 241 | } 242 | _ => { 243 | body.push(token); 244 | } 245 | } 246 | } 247 | 248 | body.extend(it); 249 | 250 | for ver in config.versions { 251 | let tag = ver.join(""); 252 | let mut ver_num = Vec::::new(); 253 | for (i, comp) in ver.iter().enumerate() { 254 | let idx = config.enums[i].iter().position(|&r| r == *comp).unwrap(); 255 | ver_num.push(idx); 256 | } 257 | out.extend(filter_versions( 258 | config, 259 | &tag, 260 | &ver_num, 261 | body.clone(), 262 | is_struct, 263 | )); 264 | } 265 | 266 | out 267 | } 268 | -------------------------------------------------------------------------------- /playground/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "playground" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | macros = { path = "../macros" } 10 | -------------------------------------------------------------------------------- /playground/src/main.rs: -------------------------------------------------------------------------------- 1 | #![feature(generic_associated_types)] 2 | 3 | use macros::versions; 4 | // use std::mem; 5 | use core::fmt::Debug; 6 | use core::fmt::Error; 7 | use core::fmt::Formatter; 8 | use core::marker::PhantomData; 9 | use core::mem::MaybeUninit; 10 | use core::num::NonZeroU64; 11 | use core::sync::atomic::{AtomicU32, Ordering}; 12 | use std::sync::Arc; 13 | // use core::ops::FnOnce; 14 | 15 | // struct GpuBox(usize, PhantomData); // Box equivalent 16 | // struct GpuRef<'a, T>(usize, PhantomData<&'a T>); // &'a T equivalent 17 | // struct GpuRefMut<'a, T>(usize, PhantomData<&'a mut T>); // &'a mut T equivalent 18 | // 19 | // struct GpuVec(usize, PhantomData); // Box<[T]> equivalent 20 | 21 | #[repr(transparent)] 22 | struct GPUPointer<'a, T>(NonZeroU64, PhantomData<&'a T>); 23 | 24 | impl<'a, T> Debug for GPUPointer<'a, T> { 25 | fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> { 26 | f.write_fmt(format_args!( 27 | "{:#x} ({})", 28 | self.0, 29 | core::any::type_name::() 30 | )) 31 | } 32 | } 33 | 34 | #[repr(transparent)] 35 | struct GPUWeakPointer(NonZeroU64, PhantomData); 36 | 37 | impl GPUWeakPointer { 38 | // The third argument is a type inference hack 39 | unsafe fn offset(&self, off: usize, _: *const U) -> GPUWeakPointer { 40 | GPUWeakPointer::(self.0.checked_add(off as u64).unwrap(), PhantomData) 41 | } 42 | } 43 | 44 | #[repr(transparent)] 45 | struct GPURawPointer(NonZeroU64); 46 | 47 | #[macro_export] 48 | macro_rules! inner_ptr { 49 | ($gpuva:expr, $($f:tt)*) => ({ 50 | fn uninit_from(_: &GPUWeakPointer) -> MaybeUninit> { 51 | core::mem::MaybeUninit::uninit() 52 | } 53 | let tmp = uninit_from($gpuva); 54 | let outer = tmp.as_ptr(); 55 | let p: *const _ = unsafe { core::ptr::addr_of!((*outer).$($f)*) }; 56 | let inner = p as *const u8; 57 | let off = unsafe { inner.offset_from(outer as *const u8) }; 58 | unsafe { $gpuva.offset(off.try_into().unwrap(), p) } 59 | }) 60 | } 61 | 62 | trait GPUStruct: 'static { 63 | type Raw<'a>: Debug; 64 | } 65 | 66 | struct GPUObject { 67 | raw: Box>, 68 | gpu_ptr: GPUWeakPointer, 69 | inner: T, 70 | } 71 | 72 | impl GPUObject { 73 | fn new(inner: T, callback: impl for<'a> FnOnce(&'a T) -> T::Raw<'a>) -> GPUObject { 74 | let r = Box::new(unsafe { std::mem::transmute_copy(&callback(&inner)) }); 75 | GPUObject:: { 76 | raw: r, 77 | gpu_ptr: GPUWeakPointer::(NonZeroU64::new(1).unwrap(), PhantomData), 78 | inner, 79 | } 80 | } 81 | 82 | fn new_prealloc( 83 | inner_cb: impl FnOnce(&GPUWeakPointer) -> T, 84 | raw_cb: impl for<'a> FnOnce(&'a T) -> T::Raw<'a>, 85 | ) -> GPUObject { 86 | let gpu_ptr = GPUWeakPointer::(NonZeroU64::new(1).unwrap(), PhantomData); 87 | let inner = inner_cb(&gpu_ptr); 88 | let r = Box::new(unsafe { std::mem::transmute_copy(&raw_cb(&inner)) }); 89 | GPUObject:: { 90 | raw: r, 91 | gpu_ptr, 92 | inner, 93 | } 94 | } 95 | 96 | fn gpu_pointer(&self) -> GPUPointer<'_, T> { 97 | GPUPointer(self.gpu_ptr.0, PhantomData) 98 | } 99 | 100 | fn with_mut( 101 | &mut self, 102 | callback: impl for<'a> FnOnce(&'a mut ::Raw<'a>, &'a mut T) -> RetVal, 103 | ) -> RetVal { 104 | unsafe { 105 | let ptr: *mut T::Raw<'static> = &mut *self.raw; 106 | callback(&mut *ptr, &mut *(&mut self.inner as *mut _)) 107 | } 108 | } 109 | 110 | fn with( 111 | &self, 112 | callback: impl for<'a> FnOnce(&'a ::Raw<'a>, &'a T) -> RetVal, 113 | ) -> RetVal { 114 | unsafe { 115 | let ptr: *const T::Raw<'static> = &*self.raw; 116 | callback(&*ptr, &*(&self.inner as *const _)) 117 | } 118 | } 119 | } 120 | 121 | impl<'a, T: GPUStruct + Debug> Debug for GPUObject { 122 | fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> { 123 | f.debug_struct(core::any::type_name::()) 124 | .field("raw", &format_args!("{:#X?}", &self.raw)) 125 | .field("inner", &format_args!("{:#X?}", &self.inner)) 126 | .finish() 127 | } 128 | } 129 | 130 | struct GPUArray { 131 | raw: Box>, 132 | gpu_ptr: NonZeroU64, 133 | } 134 | 135 | impl GPUArray { 136 | fn new(data: Vec) -> GPUArray { 137 | GPUArray:: { 138 | raw: Box::new(data), 139 | gpu_ptr: NonZeroU64::new(1).unwrap(), 140 | } 141 | } 142 | 143 | fn gpu_pointer(&self) -> GPUPointer<'_, &'_ [T]> { 144 | GPUPointer(self.gpu_ptr, PhantomData) 145 | } 146 | 147 | fn len(&self) -> usize { 148 | (*self.raw).len() 149 | } 150 | 151 | fn as_slice(&self) -> &[T] { 152 | &*self.raw 153 | } 154 | 155 | fn as_mut_slice(&mut self) -> &mut [T] { 156 | &mut *self.raw 157 | } 158 | } 159 | 160 | impl Drop for GPUObject { 161 | fn drop(&mut self) { 162 | println!("Dropping {}", core::any::type_name::(),); 163 | } 164 | } 165 | 166 | impl<'a, T: Copy + Debug> Debug for GPUArray { 167 | fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> { 168 | f.debug_struct(core::any::type_name::()) 169 | .field("array", &format_args!("{:#X?}", &self.raw)) 170 | .finish() 171 | } 172 | } 173 | 174 | #[repr(C)] 175 | #[derive(Debug, Default)] 176 | struct RawBufferManagerBlockControl<'a> { 177 | total: AtomicU32, 178 | _t: PhantomData<&'a ()>, 179 | } 180 | 181 | #[derive(Debug)] 182 | struct BufferManagerBlockControl {} 183 | 184 | impl GPUStruct for BufferManagerBlockControl { 185 | type Raw<'a> = RawBufferManagerBlockControl<'a>; 186 | } 187 | 188 | #[versions(AGX)] 189 | #[derive(Debug)] 190 | #[repr(C)] 191 | struct RawBufferManagerInfo<'a> { 192 | gpu_counter: u32, 193 | unk_4: u32, 194 | #[ver(V < V13_0b4)] 195 | unk_4_0: u32, 196 | block_count: AtomicU32, 197 | block_ctl_addr: GPUPointer<'a, BufferManagerBlockControl>, 198 | } 199 | 200 | #[versions(AGX)] 201 | #[derive(Debug)] 202 | struct BufferManagerInfo { 203 | block_control: GPUObject, 204 | } 205 | 206 | #[versions(AGX)] 207 | impl GPUStruct for BufferManagerInfo::ver { 208 | type Raw<'a> = RawBufferManagerInfo::ver<'a>; 209 | } 210 | 211 | #[versions(AGX)] 212 | #[derive(Debug)] 213 | #[repr(C)] 214 | struct RawTAJob<'a> { 215 | buffer_mgr: GPUPointer<'a, BufferManagerInfo::ver>, 216 | micro_sequence: GPUPointer<'a, &'a [u8]>, 217 | foo: u64, 218 | } 219 | 220 | #[versions(AGX)] 221 | #[derive(Debug)] 222 | struct TAJob { 223 | buffer_mgr: Arc>, 224 | micro_sequence: GPUMicroSequence, 225 | } 226 | 227 | #[versions(AGX)] 228 | impl GPUStruct for TAJob::ver { 229 | type Raw<'a> = RawTAJob::ver<'a>; 230 | } 231 | 232 | type GPUMicroSequence = GPUArray; 233 | 234 | struct GPUMicroSequenceBuilder { 235 | ops: Vec, 236 | } 237 | 238 | impl GPUMicroSequenceBuilder { 239 | fn new() -> GPUMicroSequenceBuilder { 240 | GPUMicroSequenceBuilder { ops: Vec::new() } 241 | } 242 | 243 | fn add(&mut self, op: T) { 244 | let p: *const T = &op; 245 | let p: *const u8 = p as *const u8; 246 | let s: &[u8] = unsafe { std::slice::from_raw_parts(p, std::mem::size_of::()) }; 247 | self.ops.extend(s); 248 | } 249 | 250 | fn build(self) -> GPUMicroSequence { 251 | GPUArray::::new(self.ops) 252 | } 253 | } 254 | 255 | trait GPUMSOp {} 256 | 257 | #[repr(C, packed(4))] 258 | struct SomeOp { 259 | magic: u8, 260 | foo_p: GPUWeakPointer, 261 | } 262 | 263 | impl SomeOp { 264 | const MAGIC: u8 = 4; 265 | } 266 | 267 | impl GPUMSOp for SomeOp {} 268 | 269 | #[versions(AGX)] 270 | struct GPUDriver {} 271 | 272 | #[versions(AGX)] 273 | impl GPUDriver::ver { 274 | fn run() { 275 | let mut ctl = GPUObject::new(BufferManagerBlockControl {}, |inner| { 276 | RawBufferManagerBlockControl { 277 | total: AtomicU32::new(0), 278 | ..Default::default() 279 | } 280 | }); 281 | 282 | let mut ctl2 = 283 | GPUObject::::new(BufferManagerBlockControl {}, |inner| { 284 | RawBufferManagerBlockControl { 285 | total: AtomicU32::new(0), 286 | ..Default::default() 287 | } 288 | }); 289 | 290 | dbg!(ctl2); 291 | 292 | let mut mgr = GPUObject::::new( 293 | BufferManagerInfo::ver { block_control: ctl }, 294 | |inner| RawBufferManagerInfo::ver { 295 | gpu_counter: 0, 296 | unk_4: 0, 297 | #[ver(V < V13_0b4)] 298 | unk_4_0: 1, 299 | block_count: AtomicU32::new(0), 300 | block_ctl_addr: inner.block_control.gpu_pointer(), 301 | }, 302 | ); 303 | 304 | dbg!(&mgr); 305 | 306 | mgr.with_mut(|raw, inner| { 307 | raw.gpu_counter = 2; 308 | inner.block_control.with(|raw, inner| { 309 | raw.total.fetch_add(1, Ordering::Relaxed); 310 | }); 311 | }); 312 | 313 | let arc = Arc::new(mgr); 314 | 315 | let mut ta = GPUObject::::new_prealloc( 316 | |p| { 317 | let mut ms = GPUMicroSequenceBuilder::new(); 318 | ms.add(SomeOp { 319 | magic: SomeOp::MAGIC, 320 | foo_p: inner_ptr!(p, foo), 321 | }); 322 | TAJob::ver { 323 | buffer_mgr: arc.clone(), 324 | micro_sequence: ms.build(), 325 | } 326 | }, 327 | |inner| RawTAJob::ver { 328 | buffer_mgr: inner.buffer_mgr.gpu_pointer(), 329 | micro_sequence: inner.micro_sequence.gpu_pointer(), 330 | foo: 0x1234, 331 | }, 332 | ); 333 | 334 | arc.with(|raw, inner| { 335 | inner.block_control.with(|raw, inner| { 336 | raw.total.fetch_add(1, Ordering::Relaxed); 337 | }); 338 | }); 339 | 340 | dbg!(&ta); 341 | } 342 | } 343 | 344 | fn main() { 345 | GPUDriverG13GV13_0b4::run(); 346 | GPUDriverG13GV12_3::run(); 347 | } 348 | --------------------------------------------------------------------------------