Vector Legalization lowering pass

[mcourteaux]

⭐ New feature: Wrote a vector legalization lowering pass that comes near the end of the lowering. For loop Device APIs determine the maximal lane count for the expressions inside that for. Shuffles get lifted into their own variable, such that splitting into groups of lanes is done without recalculations.

• There are a few unsupported scenarios, which are reported as internal_errors:
  1. VectorReduce with output lanes > 1.
  2. Reinterpret with input/output having different number of bits per element.
• This lowering pass was stress tested by running all tests in the test suite with artificial vector lane limit of 4 (this artificial limit is disabled again in this PR).
• The test error/metal_vector_too_large is converted into correctness/metal_long_vectors as this is now supported.

🧹 Cleanup errors: no newline needed for HeapPrinter, and helper macro vk_report_error to print error codes. This trailing newline is pretty much everywhere in the codebase with a 50% probability for internal_assert and internal_error. This could use a more broad cleanup.

---

Fixes #8628 (see for details): use OpVectorShuffle instead of OpCompositeInsert.

[mcourteaux]

Regarding the build failure: some simplification rules in Simplify_Shuffle.cpp rewrite shuffles and thereby exceed the maximal vector length on the GPU backend (triggered by my randomized shuffle indices test). I hate this, as my PR was supposed to be simple, and now I'm uncovering more bugs than I signed up for 😝

I see two possible solutions:

1. Limit the rules to never exceed the maximal vector length per backend. I don't like this for several reasons:
   • These simplification rules are embedded in the general simplify() call, which can happen at any point, and is not yet aware of selected GPU APIs.
   • This makes the simplification rules restricted, and might limit finding better simplifications if the internal rewriting logic cannot exceed some arbitrary vector size.
2. The GPU backends need to handle arbitrary vector size expressions, just like the LLVM backend does. You can .vectorize(x, 17) and it just works. This might not be easy to achieve, but I feel like this is the better solution, as it keeps the simplifier as is, and brings the scheduling possibilities in line with the LLVM-based codegen backends. This might be one shared lowering pass for GPU backends with a max-vector-size parameter that rewrites these expressions again. Not sure if that's easy to do...

Opinions @abadams @zvookin ?

[abadams]

I think the gpu backends need to handle arbitrary vector sizes. The rest of the compiler is free to make vectors of any size. A shared legalization pass for non-llvm backends that maps from Halide IR to Halide IR with narrower vectors might work, but seems a little tricky to get right for things like vectorreduce nodes. The other approach would be just changing how all ops are printed to handle small bundles of values, but this seems even nastier.

Basically I agree with your option 2.

[mcourteaux]

I'd consider this to be greatly out of scope of this bugfix PR. I guess I'll skip that test for now and make an issue?

Looking at the IR that triggers the error:

let t397 = ramp(.thread_id_x + g$8.s0.x.v17.base.s, g$8.extent.0, 4)
let t398 = concat_vectors(f0$8[t397], f1$8[t397])
let t404 = extract_element(t398, 1)
let t405 = extract_element(t398, 0)

which came from this:

g$8(g$8.s0.x) = let t374 = slice_vectors(
    concat_vectors(f0$8(g$8.s0.x, 0), f0$8(g$8.s0.x, 1), f0$8(g$8.s0.x, 2), f0$8(g$8.s0.x, 3)), 
    concat_vectors(f1$8(g$8.s0.x, 4), f1$8(g$8.s0.x, 5), f1$8(g$8.s0.x, 6), f1$8(g$8.s0.x, 7)),
    1, -1, 2)
  in (let t375 = (t374*t374)
     in (extract_element(t375, 0) + extract_element(t375, 1))
  )

It seems that the simplifier rules have done a good job simplifying it, but there is some simplifications missing, or the simplifier rules have gotten stuck in a local minimum. The two ramped loads of size 4 get concatenated into a vector of size 8, to then just take elements 0 and 1 out of it. Very inefficient, compared to just doing two loads (or one ramped load). Of course, this is due to the unnatural way of constructing it with all these explicit shuffles in the test, but perhaps, having Halide simplify this further might be achievable for this PR? @abadams Any ideas on improving the codegen for this?

[mcourteaux]

Aaarrrghghhhh llvm/llvm-project@735209c

The following Auto-Upgrade rules are used to maintain compatibility with
IR using the legacy intrinsics:

• llvm.nvvm.barrier0 --> llvm.nvvm.barrier.cta.sync.aligned.all(0)

Clearly doesn't work for us here... 😢

[alexreinking]

n.b. - the "fixes #NN" magic belongs on a single line (one per line if multiple) in either the PR description or the final commit description (in the GitHub interface, which by default concatenates all the intermediate commit messages). It only clutters the PR title.

[mcourteaux]

Specifically asking for review from @abadams as I uncovered another bug, but this time in the Deinterleaver transformation.

This transformation is implemented as a GraphIRMutator, but is not supposed to recurse fully, because it's goal is to produce lane-extracted Exprs from other Exprs, such as t340.odd_lanes. The bug in this instance was due to recursing too deeply:

let t99 = f0[ramp(0, 1, 8)]
let t100 = shuffle(t99, 0, 1, 2, 3)
let t101 = shuffle(t100, 0, 1)

When the deinterleaver extracts t101.even_lanes, it realizes that the resulting type has 1 lane (i.e., scalar). But then it recurses into t100, and incorrectly assumed t100 was of the same type as t101.

So, in my opinion, I don't think this Deinterleaver should ever recurse into shuffle arguments. But as I didn't write this transformation, nor do I know fully what it's purpose is, I am hesitant just deleting the recursion there:

Halide/src/Deinterleave.cpp

Lines 389 to 403 in 85a3b07

	// If this is extracting a single lane, try to recursively deinterleave rather
	// than leaving behind a shuffle.
	if (indices.size() == 1) {
	int index = indices.front();
	for (const auto &i : op->vectors) {
	if (index < i.type().lanes()) {
	ScopedValue<int> lane(starting_lane, index);
	return mutate(i);
	}
	index -= i.type().lanes();
	}
	internal_error << "extract_lane index out of bounds: " << Expr(op) << " " << index << "\n";
	}
	return Shuffle::make(op->vectors, indices);

I don't understand what the purpose of is of trying to extract odd and even lanes from a shuffle argument if the shuffle is actually just an element extraction.

[mcourteaux]

@derek-gerstmann already reviewed the Vulkan interleave codegen. That part didn't change. This involved the changes in:

• CodeGen_Vulkan_Dev.cpp
• vulkan_internal.cpp
• vulkan_resources.cpp

I worked together with @abadams on a Shuffle simplification bug in Simplify_Shuffle.cpp, so that part is already "reviewed" (pair programmed, you could say).

It'd be great if @abadams could look at the following for review:

• LegalizeVectors.cpp
• The small change in CSE.cpp to not lift Calls if they are not pure (with the exception for CallType::Halide, which can be lifted, according to a test case).
• My bugfix in Deinterleave.cpp

[abadams]

Why is this here rather than in the simplifier?

[mcourteaux]

Because I can't call the simplifier on the shuffle, and expect it to only touch the shuffle. I can only do simplify(...) which runs ALL of the simplifier logic. It's a bit pitty/unintuitive that Simplify_Shuffle.cpp is not accessible as is.

Also, I wasn't too sure I could add that to the general simplifier code either. I could try to merge the two procedures. Perhaps other places benefit from these simplifier rules too then.

[mcourteaux]

I'll try this tomorrow. It indeed is late.

[abadams]

I'm not sure it's a good idea to preserve only trace loads and trace stores, because those are supposed to be nested in other tracing events. Or is the idea that those other events won't see this mutator, because they're scalar?

[mcourteaux]

The test suite didn't show these to be nested anywhere. They are surrounded by other trace events, such as begin and end of a Func. AFAIK, they weren't nested. To be transparent: I have never ever used the tracing features. I was just looking at IR before and after legalization, to make sure it all seemed reasonable.

[abadams]

Sorry, by nested I meant they should execute after a begin_realization event (or whatever it's called), and before an end_realization, so it would be bad to drop those outer events.

[mcourteaux]

IIRC, I think they never are processed here. The begin_realization and end_realization trace calls are never involved in vectorized expressions. So this ExtractLanes mutator will never be ran on those IR nodes. Perhaps I should turn this into an internal_assert() to validate my idea.

[mcourteaux]

@alexreinking Maybe you can take a look at what I did for the tracing events. You're currently very into these.

[abadams]

How is this different to the deinterleave function in Deinterleave.cpp? Should they be unified?

[mcourteaux]

Hmm, perhaps. I think I didn't understand what Deinterleave was doing. And I'm not too sure I do now. Deinterleaver and Interleaver are doing a weird dance together which I didn't understand either.

[mcourteaux]

I'm comparing Deinterleaver and ExtractLanes. Can you have a look at their Load visitors? The alignment gets dropped if the starting lane is not 0. I don't understand why we wouldn't simply update the alignment, like I did in the ExtractLanes version. Is this an oversight in the Deinterleaver, or am I not understanding the rationale behind dropping it?

[abadams]

The alignment is the alignment of the first lane. I think your logic is only correct if the load is of a ramp with stride 1. If it's some gather of some complex expression, that's not the right way to update it. In the cases we can safely update it, the simplifier can reinfer it very easily, so I thought it best to just leave it to the next simplifier pass.

[abadams]

I'm still not sure why this can't use the existing Deinterleaver

[stevesuzuki-arm]

@alexreinking @stevesuzuki-arm Can you check if the failures on ARM are related to the work done in #8898 ? As in: does that PR fix the case below?

With #8898, this fail is no longer reproduced with seed=11290674455725750672 .
But with seed=10736607054766902027, LLVM crash happens: llvm/llvm-project#187557

[abadams]

I think the correctness_simplify failure is legit, but I think it's the expected output that is wrong (i.e. main), not the behavior in this PR.

[abadams]

Oh, oops, that's fixed in the other PR. Hopefully this didn't make too much of a merge conflict.

[codecov]

Codecov Report

❌ Patch coverage is 71.29506% with 215 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.05%. Comparing base (5c21c82) to head (34d94db).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/Deinterleave.cpp	68.32%	67 Missing and 41 partials ⚠️
src/LegalizeVectors.cpp	71.00%	43 Missing and 46 partials ⚠️
src/Simplify_Shuffle.cpp	87.05%	4 Missing and 7 partials ⚠️
src/IR.cpp	0.00%	0 Missing and 2 partials ⚠️
src/CodeGen_ARM.cpp	0.00%	0 Missing and 1 partial ⚠️
src/CodeGen_Hexagon.cpp	0.00%	0 Missing and 1 partial ⚠️
src/CodeGen_LLVM.cpp	50.00%	0 Missing and 1 partial ⚠️
src/Lower.cpp	66.66%	0 Missing and 1 partial ⚠️
src/Simplify_Let.cpp	80.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #8629      +/-   ##
==========================================
- Coverage   70.08%   70.05%   -0.03%     
==========================================
  Files         260      261       +1     
  Lines       79287    79900     +613     
  Branches    19327    19506     +179     
==========================================
+ Hits        55569    55975     +406     
- Misses      17923    18006      +83     
- Partials     5795     5919     +124     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.