[Containers] Add Array::Map #12692
[Containers] Add Array::Map #12692
Conversation
Previously, an in-place mutation could be applied to an array using `Array::MutateByApply`, but this couldn't be used for transformations that return a new array, or for transformations that return a new type. The commit adds `Array::Map`, which can map to any `ObjectRef` subclass. For mappings that return the same type, this is done by delegating to `Array::MutateByApply`, to take advantage of the same copy-on-write behavior.
With the new `Array::Map` functionality, many places that previously used explicit loops or `tvm::tir::MutateArray` can be cleaned.
| * \return The transformed array. | ||
| */ | ||
| template <typename F, typename U = std::invoke_result_t<F, T>> | ||
| Array<U> Map(F fmap) const { |
There was a problem hiding this comment.
I was curious if we could unify and migrate MutateByApply into Map?
There was a problem hiding this comment.
Possibly, and that would allow for avoiding copies in a few additional cases (e.g. map from T to Optional<T>, or to a superclass of T) that aren't currently handled. I'll take a quick stab at it and see if I can unify the two.
There was a problem hiding this comment.
Thank you on the suggestion, and it ended up being much cleaner that way. Both Map and MutateByApply are now implemented in terms of the same underlying helper function. The helper function applies both the mutate-in-place and copy-on-write optimizations, with if constexpr type checks to avoid attempting the optimization if they wouldn't be possible.
| } | ||
|
|
||
| // Normal path for incompatible types, or post-copy path for | ||
| // copy-on-write instances. |
There was a problem hiding this comment.
What will be left over on the copy-on-write instance? Will there be some items that are incompatible? How are those guaranteed to be at the end?
There was a problem hiding this comment.
What will be left over on the copy-on-write instance?
If we have compatible types, and we've reached this point, we've found at least one element for which the mapped.same_as(*it) check on line 776 has failed. In that case, output will contain everything in the range [arr->begin(), it). That is, output contains all elements that are identical, and the first non-identical element. it will point to the next element that should be transformed, and so the next loop over it can continue where the first loop left off.
Will there be some items that are incompatible?
It's entirely possible, either at compile-time or at runtime. For example, I could have an Array<PrimExpr> buffer_shape and map it to allowed ranges buffer_shape.Map([](PrimExpr expr) { return Range::FromMinExtent(0, expr);});, which would be incompatible and identified as such at compile-time. In that case, the if constexpr could identify that they cannot be represented by the same underlying array, and can skip the attempts to do so altogether.
If a type is incompatible at runtime, then it will also fail the mapped.same_as(*it) check on line 776. So if I have an Array<Var> being mapped to Array<PrimExpr> with var_array.Map([&](Var var) { return var.same_as(to_replace) ? replace_with : var;});, it may or may not be compatible, depending on whether to_replace shows up in the array.
How are those guaranteed to be at the end?
Incompatible items may occur at any point in the mapped output, even at the very first iteration. In that case, the commands executed in the conditional on !mapped.same_as(*it) are the same as would be executed up through the first iteration of the mapping loop.
// Same as the else branch on `compatible_types`
output = ArrayNode::CreateRepeated(arr->size(), U());
// For the first iteration, it is `arr->begin()`, so this would be an
// empty range [begin, begin), nothing is initialized, and this
// statement has no effect.
output->InitRange(0, arr->begin(), it);
// The newly mapped item is stored to the first location of the output.
output->SetItem(it - arr->begin(), std::move(mapped));
// The loop increment that would have happened
it++;
// `it` now points to the second element of the input, and we have one
// mapped element in the output. We're now ready to start the second
// loop, just at the second iteration instead of the first.Essentially, we only need to check for identical return values up until we find a single non-identical element, at which point we know that we can't avoid the copy anyways. But once we reach the first non-identical value, we don't need to repeat the function calls up to that point, because we know that everything is either identical (and can therefore be copied from the input) or is non-identical (is which case it is the first such non-identical value).
There was a problem hiding this comment.
Thanks for this explanation! Maybe it would be helpful to others as well to summarize this in the comment block on 796...
There was a problem hiding this comment.
No problem, and updated!
| } | ||
|
|
||
| return output; | ||
| } |
There was a problem hiding this comment.
General comment - Can you add a unit test to exercise the edge cases in MapHelper?
There was a problem hiding this comment.
Certainly, and thank you for pointing that out! There are some existing tests in container_test.cc, along with a large amount of usage when lowering TIR, but no tests that would specifically point to these edge cases.
There was a problem hiding this comment.
Tests added for each of the compatible types, to validate that copies are avoided, and to ensure correct fail-through behavior when a copy is required. A double thanks for requesting it, as it also caught a type conversion error that I had missed.
|
@junrushao any additional requests on this PR? Thank you! |
|
Awesome thank you @Lunderberg @janetsc @junrushao , the PR has been merged! |
* [Containers] Add Array::Map Previously, an in-place mutation could be applied to an array using `Array::MutateByApply`, but this couldn't be used for transformations that return a new array, or for transformations that return a new type. The commit adds `Array::Map`, which can map to any `ObjectRef` subclass. For mappings that return the same type, this is done by delegating to `Array::MutateByApply`, to take advantage of the same copy-on-write behavior. * [Refactor] Use Array::Map where possible With the new `Array::Map` functionality, many places that previously used explicit loops or `tvm::tir::MutateArray` can be cleaned. * Merge the Map and MutateInPlace implementations * Fix off-by-one error in MapHelper * Updated with unit tests for Array::Map conversions * Improved comments explaining the copy-on-write in MapHelper
Previously, an in-place mutation could be applied to an array using
Array::MutateByApply, but this couldn't be used for transformations that return a new array, or for transformations that return a new type. The PR addsArray::Map, which can map to anyObjectRefsubclass. For mappings that return the same type, this is done by delegating toArray::MutateByApply, to take advantage of the same copy-on-write behavior.For ease of review, this PR consists of two separate commits. The first implements
Array::Map, while the second performs several small refactors