Split a new mem module (for "memory" ops/passes/etc.), out of qptr.

examples/spv-lower-link-qptr-lift.rs

@@ -80,10 +80,10 @@ fn main() -> std::io::Result<()> {
             after_pass("", &module)?;
 
             // HACK(eddyb) this is roughly what Rust-GPU would need.
-            let layout_config = &spirt::qptr::LayoutConfig {
+            let layout_config = &spirt::mem::LayoutConfig {
                 abstract_bool_size_align: (1, 1),
                 logical_ptr_size_align: (4, 4),
-                ..spirt::qptr::LayoutConfig::VULKAN_SCALAR_LAYOUT
+                ..spirt::mem::LayoutConfig::VULKAN_SCALAR_LAYOUT
             };
 
             eprint_duration(|| {
@@ -92,9 +92,11 @@ fn main() -> std::io::Result<()> {
             eprintln!("qptr::lower_from_spv_ptrs");
             after_pass("qptr::lower_from_spv_ptrs", &module)?;
 
-            eprint_duration(|| spirt::passes::qptr::analyze_uses(&mut module, layout_config));
-            eprintln!("qptr::analyze_uses");
-            after_pass("qptr::analyze_uses", &module)?;
+            eprint_duration(|| {
+                spirt::passes::qptr::analyze_mem_accesses(&mut module, layout_config)
+            });
+            eprintln!("mem::analyze_accesses");
+            after_pass("mem::analyze_accesses", &module)?;
 
             eprint_duration(|| spirt::passes::qptr::lift_to_spv_ptrs(&mut module, layout_config));
             eprintln!("qptr::lift_to_spv_ptrs");

src/lib.rs

@@ -169,6 +169,7 @@ pub mod passes {
     pub mod link;
     pub mod qptr;
 }
+pub mod mem;
 pub mod qptr;
 pub mod spv;
 
@@ -397,6 +398,10 @@ pub enum Attr {
     // of `AttrSetDef::{dbg_src_loc,set_dbg_src_loc}`.
     DbgSrcLoc(OrdAssertEq<DbgSrcLoc>),
 
+    /// Memory-specific attributes (see [`mem::MemAttr`]).
+    #[from]
+    Mem(mem::MemAttr),
+
     /// `QPtr`-specific attributes (see [`qptr::QPtrAttr`]).
     #[from]
     QPtr(qptr::QPtrAttr),
@@ -489,7 +494,7 @@ pub enum DiagMsgPart {
     Attrs(AttrSet),
     Type(Type),
     Const(Const),
-    QPtrUsage(qptr::QPtrUsage),
+    MemAccesses(mem::MemAccesses),
 }
 
 /// Wrapper to limit `Ord` for interned index types (e.g. [`InternedStr`])
@@ -638,7 +643,7 @@ pub struct GlobalVarDecl {
 
     /// When `type_of_ptr_to` is `QPtr`, `shape` must be used to describe the
     /// global variable (see `GlobalVarShape`'s documentation for more details).
-    pub shape: Option<qptr::shapes::GlobalVarShape>,
+    pub shape: Option<mem::shapes::GlobalVarShape>,
 
     /// The address space the global variable will be allocated into.
     pub addr_space: AddrSpace,
@@ -947,6 +952,10 @@ pub enum DataInstKind {
     // to avoid needing special handling for recursion where it's impossible.
     FuncCall(Func),
 
+    /// Memory-specific operations (see [`mem::MemOp`]).
+    #[from]
+    Mem(mem::MemOp),
+
     /// `QPtr`-specific operations (see [`qptr::QPtrOp`]).
     #[from]
     QPtr(qptr::QPtrOp),

src/qptr/layout.rs → src/mem/layout.rs

@@ -1,6 +1,6 @@
 // FIXME(eddyb) layouts are a bit tricky: this recomputes them from several passes.
 
-use crate::qptr::shapes;
+use crate::mem::shapes;
 use crate::{
     AddrSpace, Attr, Const, ConstKind, Context, Diag, FxIndexMap, Type, TypeKind, TypeOrConst, spv,
 };
@@ -61,10 +61,10 @@ impl LayoutConfig {
         Self { min_aggregate_legacy_align: 16, ..Self::VULKAN_STANDARD_LAYOUT };
 }
 
-pub(super) struct LayoutError(pub(super) Diag);
+pub(crate) struct LayoutError(pub(crate) Diag);
 
 #[derive(Clone)]
-pub(super) enum TypeLayout {
+pub(crate) enum TypeLayout {
     Handle(HandleLayout),
     HandleArray(HandleLayout, Option<NonZeroU32>),
 
@@ -73,16 +73,16 @@ pub(super) enum TypeLayout {
 }
 
 // NOTE(eddyb) `Handle` is parameterized over the `Buffer` layout.
-pub(super) type HandleLayout = shapes::Handle<Rc<MemTypeLayout>>;
+pub(crate) type HandleLayout = shapes::Handle<Rc<MemTypeLayout>>;
 
-pub(super) struct MemTypeLayout {
-    pub(super) original_type: Type,
-    pub(super) mem_layout: shapes::MaybeDynMemLayout,
-    pub(super) components: Components,
+pub(crate) struct MemTypeLayout {
+    pub(crate) original_type: Type,
+    pub(crate) mem_layout: shapes::MaybeDynMemLayout,
+    pub(crate) components: Components,
 }
 
 // FIXME(eddyb) use proper newtypes for byte sizes.
-pub(super) enum Components {
+pub(crate) enum Components {
     Scalar,
 
     /// Vector and array elements (all of them having the same `elem` layout).
@@ -106,7 +106,7 @@ impl Components {
     /// this can return multiple components, with at most one ever being non-ZST.
     //
     // FIXME(eddyb) be more aggressive in pruning ZSTs so this can be simpler.
-    pub(super) fn find_components_containing(
+    pub(crate) fn find_components_containing(
         &self,
         // FIXME(eddyb) consider renaming such offset ranges to "extent".
         offset_range: Range<u32>,
@@ -168,7 +168,7 @@ impl Components {
 }
 
 /// Context for computing `TypeLayout`s from `Type`s (with caching).
-pub(super) struct LayoutCache<'a> {
+pub(crate) struct LayoutCache<'a> {
     cx: Rc<Context>,
     wk: &'static spv::spec::WellKnown,
 
@@ -178,7 +178,7 @@ pub(super) struct LayoutCache<'a> {
 }
 
 impl<'a> LayoutCache<'a> {
-    pub(super) fn new(cx: Rc<Context>, config: &'a LayoutConfig) -> Self {
+    pub(crate) fn new(cx: Rc<Context>, config: &'a LayoutConfig) -> Self {
         Self { cx, wk: &spv::spec::Spec::get().well_known, config, cache: Default::default() }
     }
 
@@ -197,7 +197,7 @@ impl<'a> LayoutCache<'a> {
     }
 
     /// Attempt to compute a `TypeLayout` for a given (SPIR-V) `Type`.
-    pub(super) fn layout_of(&self, ty: Type) -> Result<TypeLayout, LayoutError> {
+    pub(crate) fn layout_of(&self, ty: Type) -> Result<TypeLayout, LayoutError> {
         if let Some(cached) = self.cache.borrow().get(&ty).cloned() {
             return Ok(cached);
         }
@@ -348,7 +348,7 @@ impl<'a> LayoutCache<'a> {
         // FIXME(eddyb) !!! what if... types had a min/max size and then...
         // that would allow surrounding offsets to limit their size... but... ugh...
         // ugh this doesn't make any sense. maybe if the front-end specifies
-        // offsets with "abstract types", it must configure `qptr::layout`?
+        // offsets with "abstract types", it must configure `mem::layout`?
         let layout = if spv_inst.opcode == wk.OpTypeBool {
             // FIXME(eddyb) make this properly abstract instead of only configurable.
             scalar_with_size_and_align(self.config.abstract_bool_size_align)

src/mem/mod.rs

@@ -0,0 +1,151 @@
+//! Memory operations, analyses and transformations.
+//
+// FIXME(eddyb) document at least these aspects:
+// - "memory" = indirect storage of data and/or resources
+// - (untyped) "data" = mix of plain bytes and pointers (as per RalfJ blog post)
+//   (does "non-data memory" actually make sense? could be considered typed?)
+//
+// FIXME(eddyb) dig into past notes (e.g. `qptr::legalize`, Rust-GPU release notes,
+// https://github.com/EmbarkStudios/spirt/pull/24, etc.) for useful docs.
+//
+// FIXME(eddyb) consider taking this into a more (R)VSDG "state type" direction.
+
+use crate::{OrdAssertEq, Type};
+use std::collections::BTreeMap;
+use std::num::NonZeroU32;
+use std::rc::Rc;
+
+// NOTE(eddyb) all the modules are declared here, but they're documented "inside"
+// (i.e. using inner doc comments).
+pub mod analyze;
+// FIXME(eddyb) make this public?
+pub(crate) mod layout;
+pub mod shapes;
+
+pub use layout::LayoutConfig;
+
+/// Memory-specific attributes ([`Attr::Mem`]).
+#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub enum MemAttr {
+    /// When applied to a `GlobalVar` or `FuncLocalVar`, this tracks all possible
+    /// access patterns its memory may be subjected to (see [`MemAccesses`]).
+    Accesses(OrdAssertEq<MemAccesses>),
+}
+
+#[derive(Clone, PartialEq, Eq, Hash)]
+pub enum MemAccesses {
+    /// Accesses to one or more handles (i.e. optionally indexed by
+    /// [`crate::qptr::QPtrOp::HandleArrayIndex`]), which can be:
+    /// - `Handle::Opaque(handle_type)`: all accesses involve [`MemOp::Load`] or
+    ///   [`crate::qptr::QPtrAttr::ToSpvPtrInput`], with the common type `handle_type`
+    /// - `Handle::Buffer(data_happ)`: carries with it `data_happ`,
+    ///   i.e. the access patterns for the memory that is reached through
+    ///   [`crate::qptr::QPtrOp::BufferData`]
+    Handles(shapes::Handle<DataHapp>),
+
+    Data(DataHapp),
+}
+
+/// Data HAPP ("Hierarchical Access Pattern Partitioning"): all access patterns
+/// for some memory, structured by disjoint offset ranges ("partitions").
+///
+/// This is the core of "type recovery" (inferring typed memory from untyped),
+/// with "struct/"array" equivalents (i.e. as `DataHappKind` variants), but
+/// while it can be mapped to an explicitly laid out data type, it also tracks
+/// distinctions only needed during merging (e.g. [`DataHappKind::StrictlyTyped`]).
+///
+/// **Note**: the only reason for the recursive/"hierarchical" aspect is that
+/// (array-like) dynamic indexing allows for compact representation of repeated
+/// patterns, which can non-trivially nest for 3+ levels without losing the need
+/// for efficient representation - in fact, one can construct a worst-case like:
+/// ```ignore
+/// [(A, [(B, [(C, [(D, [T; N], X); 2], Y); 2], Z); 2], W); 2]
+/// ```
+/// (with only `N * 2**4` leaf `T`s because of the 4 `[(_, ..., _); 2]` levels,
+/// and the potential for dozens of such levels while remaining a plausible size)
+//
+// FIXME(eddyb) reconsider the name (acronym was picked to avoid harder decisions).
+#[derive(Clone, PartialEq, Eq, Hash)]
+pub struct DataHapp {
+    /// If present, this is a worst-case upper bound on the offsets at which
+    /// accesses may be perfomed.
+    //
+    // FIXME(eddyb) use proper newtypes for byte amounts.
+    //
+    // FIXME(eddyb) suboptimal naming choice, but other options are too verbose,
+    // including maybe using `RangeTo<_>` to explicitly indicate "exclusive".
+    //
+    // FIXME(eddyb) consider renaming such information to "extent", but that might
+    // be ambiguous with an offset range (as opposed to using the maximum of all
+    // *possible* `offset_range.end` values to describe a "maximum size").
+    pub max_size: Option<u32>,
+
+    pub kind: DataHappKind,
+}
+
+impl DataHapp {
+    pub const DEAD: Self = Self { max_size: Some(0), kind: DataHappKind::Dead };
+}
+
+#[derive(Clone, PartialEq, Eq, Hash)]
+pub enum DataHappKind {
+    /// Not actually accessed (only an intermediary state, during access analysis).
+    //
+    // FIXME(eddyb) use `Option<DataHapp>` instead? or an empty `Disjoint`?
+    Dead,
+
+    // FIXME(eddyb) replace the two leaves with e.g. `Leaf(Type, LeafAccessKind)`.
+    //
+    //
+    /// Accesses through typed pointers (e.g. via unknown SPIR-V instructions),
+    /// requiring a specific choice of pointee type which cannot be modified,
+    /// and has to be reused as-is, when lifting to typed memory.
+    ///
+    /// Other overlapping accesses can be merged into this one as long as they
+    /// can be fully expressed using the (transitive) components of this type.
+    StrictlyTyped(Type),
+
+    /// Direct accesses (e.g. [`MemOp::Load`], [`MemOp::Store`]), which can be
+    /// decomposed as necessary (down to individual scalar leaves), to allow
+    /// maximal merging opportunities.
+    //
+    // FIXME(eddyb) track whether accesses are `Load`s and/or `Store`s, to allow
+    // inferring `NonWritable`/`NonReadable` annotations, as well.
+    Direct(Type),
+
+    /// Partitioning into disjoint offset ranges (the map is keyed by the start
+    /// of the offset range, while the end is implied by its corresponding value),
+    /// requiring a "struct" type, when lifting to typed memory.
+    //
+    // FIXME(eddyb) make this non-nestable and the fundamental basis of "HAPP".
+    Disjoint(Rc<BTreeMap<u32, DataHapp>>),
+
+    /// `Disjoint` counterpart for dynamic offsetting, requiring an "array" type,
+    /// when lifting to typed memory, with one single element type being repeated
+    /// across the entire size, at all offsets that are a multiple of `stride`.
+    Repeated {
+        // FIXME(eddyb) this feels inefficient.
+        element: Rc<DataHapp>,
+        stride: NonZeroU32,
+    },
+}
+
+/// Memory-specific operations ([`DataInstKind::Mem`]).
+#[derive(Clone, PartialEq, Eq, Hash)]
+pub enum MemOp {
+    // HACK(eddyb) `OpVariable` replacement, which itself should not be kept as
+    // a `SpvInst` - once fn-local variables are lowered, this should go there.
+    FuncLocalVar(shapes::MemLayout),
+
+    /// Read a single value from a pointer (`inputs[0]`).
+    //
+    // FIXME(eddyb) limit this to data - and scalars, maybe vectors at most.
+    Load,
+
+    /// Write a single value (`inputs[1]`) to a pointer (`inputs[0]`).
+    //
+    // FIXME(eddyb) limit this to data - and scalars, maybe vectors at most.
+    Store,
+    //
+    // FIXME(eddyb) implement more ops (e.g. copies, atomics).
+}

src/qptr/shapes.rs → src/mem/shapes.rs

@@ -1,6 +1,7 @@
 //! Variable shapes (untyped memory layouts vs abstract resources).
 //
 // FIXME(eddyb) does this need its own module still?
+// TODO(eddyb) strongly consider moving these to `mem/mod.rs`.
 
 use crate::{AddrSpace, Type};
 use std::num::NonZeroU32;
@@ -22,6 +23,7 @@ pub enum GlobalVarShape {
     },
 
     // FIXME(eddyb) unify terminology around "concrete"/"memory"/"untyped (data)".
+    // TODO(eddyb) strongly consider renaming this to just `Data`.
     UntypedData(MemLayout),
 
     /// Non-memory pipeline interface, which must keep the exact original type,
@@ -73,7 +75,7 @@ pub enum Handle<BL = MaybeDynMemLayout> {
     // instead of being treated like a buffer?
     //
     // FIXME(eddyb) should this be a `Type` of its own, that can be loaded from
-    // a handle `QPtr`, and then has data pointer / length ops *on that*?
+    // a handle pointer, and then has data pointer / length ops *on that*?
     Buffer(AddrSpace, BL),
 }
 
@@ -83,11 +85,12 @@ pub enum Handle<BL = MaybeDynMemLayout> {
 /// and are both kept track of to detect ambiguity in implicit layouts, e.g.
 /// field offsets when the `Offset` decoration isn't being used.
 /// Note, however, that `legacy_align` can be raised to "extended" alignment,
-/// or completeley ignored, using [`LayoutConfig`](crate::qptr::LayoutConfig).
+/// or completeley ignored, using [`LayoutConfig`](crate::mem::LayoutConfig).
 ///
 /// Only `align` is *required*, that is `size % align == 0` must be always enforced.
 //
 // FIXME(eddyb) consider supporting specialization-constant-length arrays.
+// TODO(eddyb) strongly consider renaming this to `DataLayout`.
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
 pub struct MemLayout {
     // FIXME(eddyb) use proper newtypes (and log2 for align!).

src/passes/qptr.rs

@@ -4,7 +4,7 @@ use crate::qptr;
 use crate::visit::{InnerVisit, Visitor};
 use crate::{AttrSet, Const, Context, Func, FxIndexSet, GlobalVar, Module, Type};
 
-pub fn lower_from_spv_ptrs(module: &mut Module, layout_config: &qptr::LayoutConfig) {
+pub fn lower_from_spv_ptrs(module: &mut Module, layout_config: &crate::mem::LayoutConfig) {
     let cx = &module.cx();
 
     let (seen_global_vars, seen_funcs) = {
@@ -34,11 +34,13 @@ pub fn lower_from_spv_ptrs(module: &mut Module, layout_config: &qptr::LayoutConf
     }
 }
 
-pub fn analyze_uses(module: &mut Module, layout_config: &qptr::LayoutConfig) {
-    qptr::analyze::InferUsage::new(module.cx(), layout_config).infer_usage_in_module(module);
+// FIXME(eddyb) this doesn't really belong in `qptr`.
+pub fn analyze_mem_accesses(module: &mut Module, layout_config: &crate::mem::LayoutConfig) {
+    crate::mem::analyze::GatherAccesses::new(module.cx(), layout_config)
+        .gather_accesses_in_module(module);
 }
 
-pub fn lift_to_spv_ptrs(module: &mut Module, layout_config: &qptr::LayoutConfig) {
+pub fn lift_to_spv_ptrs(module: &mut Module, layout_config: &crate::mem::LayoutConfig) {
     let cx = &module.cx();
 
     let (seen_global_vars, seen_funcs) = {