From b5bdf121acb3df37ff66e64a79a59b14db7cbf4d Mon Sep 17 00:00:00 2001 From: Manuel Saelices Date: Sun, 2 Mar 2025 22:57:25 +0100 Subject: [PATCH 1/4] Optimize compiler times by prevent huge loop unrolling when filling inline arrays. --- mojo/stdlib/src/collections/inline_array.mojo | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/mojo/stdlib/src/collections/inline_array.mojo b/mojo/stdlib/src/collections/inline_array.mojo index 3cb3161e53..817109a2fd 100644 --- a/mojo/stdlib/src/collections/inline_array.mojo +++ b/mojo/stdlib/src/collections/inline_array.mojo @@ -140,8 +140,21 @@ struct InlineArray[ _inline_array_construction_checks[size]() __mlir_op.`lit.ownership.mark_initialized`(__get_mvalue_as_litref(self)) + alias unroll_threshold = 1000 + alias unrolled = size // unroll_threshold + alias remainder = size % unroll_threshold + @parameter - for i in range(size): + for i in range(unrolled): + for j in range(unroll_threshold): + var ptr = UnsafePointer.address_of( + self.unsafe_get(i * unroll_threshold + j) + ) + ptr.init_pointee_copy(fill) + + # Fill the remainder + @parameter + for i in range(unrolled * unroll_threshold, size): var ptr = UnsafePointer.address_of(self.unsafe_get(i)) ptr.init_pointee_copy(fill) From bb401c4700a294a463bd82ab598ee67f465a1d4a Mon Sep 17 00:00:00 2001 From: Manuel Saelices Date: Mon, 3 Mar 2025 21:24:47 +0100 Subject: [PATCH 2/4] Use math.align_down() that makes the intent more obvious Unroll the inner loop Signed-off-by: Manuel Saelices --- mojo/stdlib/src/collections/inline_array.mojo | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/mojo/stdlib/src/collections/inline_array.mojo b/mojo/stdlib/src/collections/inline_array.mojo index 817109a2fd..7131816840 100644 --- a/mojo/stdlib/src/collections/inline_array.mojo +++ b/mojo/stdlib/src/collections/inline_array.mojo @@ -22,6 +22,7 @@ from collections import InlineArray from collections._index_normalization import normalize_index from sys.intrinsics import _type_is_eq +import math from memory import UnsafePointer from memory.maybe_uninitialized import UnsafeMaybeUninitialized @@ -140,21 +141,21 @@ struct InlineArray[ _inline_array_construction_checks[size]() __mlir_op.`lit.ownership.mark_initialized`(__get_mvalue_as_litref(self)) - alias unroll_threshold = 1000 - alias unrolled = size // unroll_threshold - alias remainder = size % unroll_threshold + alias batch_size = 1000 + alias unroll_end = math.align_down(size, batch_size) - @parameter - for i in range(unrolled): - for j in range(unroll_threshold): + for i in range(0, unroll_end, batch_size): + + @parameter + for j in range(batch_size): var ptr = UnsafePointer.address_of( - self.unsafe_get(i * unroll_threshold + j) + self.unsafe_get(i * batch_size + j) ) ptr.init_pointee_copy(fill) # Fill the remainder @parameter - for i in range(unrolled * unroll_threshold, size): + for i in range(unroll_end, size): var ptr = UnsafePointer.address_of(self.unsafe_get(i)) ptr.init_pointee_copy(fill) From edadb9814bbb5d1ae7ab3ac8c22f2c6cfd4e88e9 Mon Sep 17 00:00:00 2001 From: Manuel Saelices Date: Mon, 3 Mar 2025 21:26:20 +0100 Subject: [PATCH 3/4] Smaller batch_size in the InlineArray(fill) constructor, which is now a parameter Signed-off-by: Manuel Saelices --- mojo/stdlib/src/collections/inline_array.mojo | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mojo/stdlib/src/collections/inline_array.mojo b/mojo/stdlib/src/collections/inline_array.mojo index 7131816840..dd2cc203ac 100644 --- a/mojo/stdlib/src/collections/inline_array.mojo +++ b/mojo/stdlib/src/collections/inline_array.mojo @@ -132,7 +132,7 @@ struct InlineArray[ @always_inline @implicit - fn __init__(out self, fill: Self.ElementType): + fn __init__[batch_size: Int = 100](out self, fill: Self.ElementType): """Constructs an empty array where each element is the supplied `fill`. Args: @@ -141,7 +141,6 @@ struct InlineArray[ _inline_array_construction_checks[size]() __mlir_op.`lit.ownership.mark_initialized`(__get_mvalue_as_litref(self)) - alias batch_size = 1000 alias unroll_end = math.align_down(size, batch_size) for i in range(0, unroll_end, batch_size): From 212caea18f10592b61de038ff73df46b9af496ec Mon Sep 17 00:00:00 2001 From: Manuel Saelices Date: Mon, 3 Mar 2025 22:46:23 +0100 Subject: [PATCH 4/4] Change the batch_size to be power of two. Add param to the docstring Signed-off-by: Manuel Saelices --- mojo/stdlib/src/collections/inline_array.mojo | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/mojo/stdlib/src/collections/inline_array.mojo b/mojo/stdlib/src/collections/inline_array.mojo index dd2cc203ac..3e259365ab 100644 --- a/mojo/stdlib/src/collections/inline_array.mojo +++ b/mojo/stdlib/src/collections/inline_array.mojo @@ -132,9 +132,12 @@ struct InlineArray[ @always_inline @implicit - fn __init__[batch_size: Int = 100](out self, fill: Self.ElementType): + fn __init__[batch_size: Int = 64](out self, fill: Self.ElementType): """Constructs an empty array where each element is the supplied `fill`. + Parameters: + batch_size: The number of elements to unroll for filling the array. + Args: fill: The element to fill each index. """