Support lgalloc for columnar #31230

antiguru · 2025-01-29T14:58:14Z

Adds support for allocating memory for columnar's aligned type from
lgalloc. Adds a flag enable_lgalloc_columnar to control the behavior.

Checklist

This PR has adequate test coverage / QA involvement has been duly considered. (trigger-ci for additional test/nightly runs)
This PR has an associated up-to-date design doc, is a design doc (template), or is sufficiently small to not require a design.
If this PR evolves an existing $T ⇔ Proto$T mapping (possibly in a backwards-incompatible way), then it is tagged with a T-proto label.
If this PR will require changes to cloud orchestration or tests, there is a companion cloud PR to account for those changes that is tagged with the release-blocker label (example).
If this PR includes major user-facing behavior changes, I have pinged the relevant PM to schedule a changelog post.

Adds support for allocating memory for columnar's aligned type from lgalloc. Adds a flag `enable_lgalloc_columnar` to control the behavior. Signed-off-by: Moritz Hoffmann <[email protected]> # Conflicts: # src/compute/src/logging/differential.rs # src/compute/src/logging/initialize.rs # src/compute/src/logging/reachability.rs Signed-off-by: Moritz Hoffmann <[email protected]>

ParkMyCar · 2025-01-30T21:50:39Z

src/timely-util/src/containers.rs

+    /// Allocate a region of memory with a capacity of at least `len` that is aligned to 8 bytes
+    /// and zeroed.
+    #[inline]
+    pub(crate) fn alloc_aligned_zeroed(len: usize) -> Region<u64> {
+        if enable_lgalloc_columnar() {
+            Region::new_auto_zeroed(len)
+        } else {
+            Region::new_heap_zeroed(len)
+        }
+    }


Mind adding to this comment that it's aligned to 8 bytes because it's a Region of u64s?

ParkMyCar · 2025-01-30T21:56:32Z

src/timely-util/src/containers.rs

+                        assert!(align <= 8);
+                        let length = u64::cast_from(bytes.len());
+                        let length_slice = bytemuck::cast_slice(std::slice::from_ref(&length));
+                        writer.write_all(length_slice).unwrap();
+                        writer.write_all(bytes).unwrap();
+                        let padding = usize::cast_from((8 - (length % 8)) % 8);
+                        writer.write_all(&[0; 8][..padding]).unwrap();


I'm not totally up-to-date on the encoded layout of a Column, but my read is this layouts data like:

<u64 encoded in native endianness><buffer whose length is the u64><optional padding>

If that's what it's supposed to do then it LGTM!

teskje · 2025-01-31T12:34:52Z

src/compute-types/src/dyncfgs.rs

+pub const ENABLE_LGALLOC_COLUMNAR: Config<bool> = Config::new(
+    "enable_lgalloc_columnar",


Nit: Call this enable_columnar_lgalloc for consistency with the columniation flag?

teskje · 2025-01-31T12:35:09Z

src/compute-types/src/dyncfgs.rs

+pub const ENABLE_LGALLOC_COLUMNAR: Config<bool> = Config::new(
+    "enable_lgalloc_columnar",
+    true,
+    "Enable allocating alignd regions in columnar from lgalloc.",


Suggested change

"Enable allocating alignd regions in columnar from lgalloc.",

"Enable allocating aligned regions in columnar from lgalloc.",

teskje · 2025-01-31T12:47:02Z

src/ore/src/region.rs

+            }
+        }
+        // Fall-through
+        Region::Heap(vec![0; capacity])


Suggested change

Region::Heap(vec![0; capacity])

Self::new_heap_zeroed(capacity)

teskje · 2025-01-31T12:49:02Z

src/timely-util/src/containers.rs

    use timely::bytes::arc::Bytes;
    use timely::container::PushInto;
    use timely::dataflow::channels::ContainerBytes;
    use timely::Container;

+    use crate::containers::alloc_aligned_zeroed;


Isn't this import redundant, since it's importing from the current module?

teskje · 2025-01-31T12:51:10Z

src/ore/src/region.rs

+    pub fn new_mmap_zeroed(capacity: usize) -> Result<Region<u64>, lgalloc::AllocError> {
+        lgalloc::allocate::<u64>(capacity).map(|(ptr, capacity, handle)| {
+            // SAFETY: `allocate` returns a valid memory block
+            unsafe { ptr.as_ptr().write_bytes(0, capacity) }


This seems like something that could introduce performance regressions once we start switching more things from columnation to columnar. Is the zeroing required for columnar to work?

teskje · 2025-01-31T12:56:11Z

src/timely-util/src/containers.rs

    use timely::container::PushInto;
    use timely::container::{ContainerBuilder, LengthPreservingContainerBuilder};

    use super::Column;

+    thread_local! {
+        static ENABLE_LGALLOC_COLUMNAR: std::cell::RefCell<bool> = const { std::cell::RefCell::new(false) };


Suggested change

static ENABLE_LGALLOC_COLUMNAR: std::cell::RefCell<bool> = const { std::cell::RefCell::new(false) };

static ENABLE_LGALLOC_COLUMNAR: std::cell::RefCell<bool> = std::cell::RefCell::new(false);

You could also use an AtomicBool here, to rule out the possibility of panics.

teskje · 2025-01-31T13:02:22Z

src/timely-util/src/containers.rs

+                        let length = u64::cast_from(bytes.len());
+                        let length_slice = bytemuck::cast_slice(std::slice::from_ref(&length));
+                        writer.write_all(length_slice).unwrap();
+                        writer.write_all(bytes).unwrap();
+                        let padding = usize::cast_from((8 - (length % 8)) % 8);
+                        writer.write_all(&[0; 8][..padding]).unwrap();


Is there a columnar method we could invoke so we don't have to worry about this code diverging from the columnar implementation?

teskje · 2025-01-31T13:09:21Z

src/timely-util/src/containers/array.rs

+        // SAFETY: casting `elements` to a `*const [T]` is safe since the caller guarantees that
+        // `slice` is initialized, and `MaybeUninit` is guaranteed to have the same layout as `T`.
+        // The pointer obtained is valid since it refers to memory owned by `elements` which is a
+        // reference and thus guaranteed to be valid for reads.


This safety argument talks about *const and reads, but we're casting to a *mut and deference that for writes. I think a similar argument still holds, based on that elements is mutably borrowed here, so we can produce a &mut into it.

frankmcsherry

Seems good, though there are a few moments I'm not entirely clear on. But structurally the PR seems ok!

frankmcsherry · 2025-01-30T19:46:08Z

src/ore/src/region.rs

@@ -342,6 +344,48 @@ impl<T> Region<T> {
    }
 }

+impl Region<u64> {


Does this exist because of columnar's u64 bias? If so, is this great news and we love it, or is this a work around and we're not happy? Mostly trying to grok what's going on here and for what reason, which I imagine is clearer to you than it is to me.

frankmcsherry · 2025-01-31T15:52:31Z

src/timely-util/src/containers/array.rs

@@ -120,6 +128,22 @@ impl<T> Deref for Array<T> {
    }
 }

+impl<T> DerefMut for Array<T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        // TODO: Use `slice_assume_init_ref` once stable.


Probably slice_assume_init_mut rather than _ref, I think?

If nothing else, we'll want to have the unsafe justification for _mut instead, I think. Maybe a nit, but "have the right explanation when we use unsafe" seems like a great place to pick nits.

antiguru requested review from a team and benesch as code owners January 29, 2025 14:58

antiguru requested review from teskje and frankmcsherry January 29, 2025 14:58

antiguru force-pushed the columnar_lgalloc branch from 6355228 to 2fb406a Compare January 29, 2025 16:06

ParkMyCar approved these changes Jan 30, 2025

View reviewed changes

teskje reviewed Jan 31, 2025

View reviewed changes

frankmcsherry approved these changes Jan 31, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support lgalloc for columnar #31230

Support lgalloc for columnar #31230

antiguru commented Jan 29, 2025

ParkMyCar Jan 30, 2025

ParkMyCar Jan 30, 2025

teskje Jan 31, 2025

teskje Jan 31, 2025

teskje Jan 31, 2025

teskje Jan 31, 2025

teskje Jan 31, 2025

teskje Jan 31, 2025

teskje Jan 31, 2025

teskje Jan 31, 2025

teskje Jan 31, 2025

frankmcsherry left a comment

frankmcsherry Jan 30, 2025

frankmcsherry Jan 31, 2025

frankmcsherry Jan 31, 2025

		pub const ENABLE_LGALLOC_COLUMNAR: Config<bool> = Config::new(
		"enable_lgalloc_columnar",

	"Enable allocating alignd regions in columnar from lgalloc.",
	"Enable allocating aligned regions in columnar from lgalloc.",

	Region::Heap(vec![0; capacity])
	Self::new_heap_zeroed(capacity)

	static ENABLE_LGALLOC_COLUMNAR: std::cell::RefCell<bool> = const { std::cell::RefCell::new(false) };
	static ENABLE_LGALLOC_COLUMNAR: std::cell::RefCell<bool> = std::cell::RefCell::new(false);

Support lgalloc for columnar #31230

Are you sure you want to change the base?

Support lgalloc for columnar #31230

Conversation

antiguru commented Jan 29, 2025

Checklist

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

frankmcsherry left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment