Skip to content

[mlir][spirv] Add conversion pass to rewrite splat constant composite… #148910

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 12 commits into from
Jul 18, 2025
Merged

[mlir][spirv] Add conversion pass to rewrite splat constant composite… #148910

Changes from 1 commit
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
936435e
[mlir][spirv] Add conversion pass to rewrite splat constant composite…
mahabadm Jul 2, 2025
14b70ce
Minor bug fix and revision
mahabadm Jul 16, 2025
bee1e45
Merge branch 'main' into constant_composite_replicates_conversion_pass
mahabadm Jul 16, 2025
51ce7ac
Addressing code review comments
mahabadm Jul 16, 2025
d065a5a
Addressing further code review comments
mahabadm Jul 17, 2025
b718e34
Addressing further code review comments
mahabadm Jul 17, 2025
8234e68
Minor typo fix
mahabadm Jul 17, 2025
568bccb
Addressing further code review comments
mahabadm Jul 17, 2025
ae1eb18
Addressing further code review comments
mahabadm Jul 17, 2025
6ea73e3
Slight change of logic for value type detection
mahabadm Jul 18, 2025
a368eb7
Merge branch 'main' into constant_composite_replicates_conversion_pass
mahabadm Jul 18, 2025
1584f58
Addressing further code review comments
mahabadm Jul 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 11 additions & 14 deletions mlir/lib/Dialect/SPIRV/Transforms/ConvertToReplicatedConstantCompositePass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,27 +25,24 @@ namespace {
static std::pair<Attribute, uint32_t>
getSplatAttrAndNumElements(Attribute valueAttr) {
Attribute attr;
uint32_t splatCount = 0;
uint32_t numElements = 0;
if (auto splatAttr = dyn_cast<SplatElementsAttr>(valueAttr)) {
return {splatAttr.getSplatValue<Attribute>(), splatAttr.size()};
}
if (auto arrayAttr = dyn_cast<ArrayAttr>(valueAttr)) {
if (llvm::all_equal(arrayAttr)) {
attr = arrayAttr[0];
splatCount = arrayAttr.size();
}
numElements = arrayAttr.size();

if (attr) {
// Find the inner-most splat value for array of composites
std::pair<Attribute, uint32_t> newSplatAttrAndCount =
getSplatAttrAndNumElements(attr);
if (newSplatAttrAndCount.first) {
return newSplatAttrAndCount;
auto [newAttr, newNumElements] = getSplatAttrAndNumElements(attr);
if (newAttr) {
return {newAttr, numElements * newNumElements};
}
}
}

return {attr, splatCount};
return {attr, numElements};
}

struct ConstantOpConversion final : OpRewritePattern<spirv::ConstantOp> {
Expand All @@ -57,16 +54,16 @@ struct ConstantOpConversion final : OpRewritePattern<spirv::ConstantOp> {
if (!compositeType)
return rewriter.notifyMatchFailure(op, "not a composite constant");

auto [splatAttr, splatCount] = getSplatAttrAndNumElements(op.getValue());
if (!splatAttr)
auto [attr, numElements] = getSplatAttrAndNumElements(op.getValue());
if (!attr)
return rewriter.notifyMatchFailure(op, "composite is not splat");

if (splatCount == 1)
if (numElements == 1)
return rewriter.notifyMatchFailure(op,
"composite has only one constituent");

rewriter.replaceOpWithNewOp<spirv::EXTConstantCompositeReplicateOp>(
op, op.getType(), splatAttr);
op, op.getType(), attr);
return success();
}
};
Expand All @@ -86,7 +83,7 @@ struct SpecConstantCompositeOpConversion final
return rewriter.notifyMatchFailure(op,
"composite has only one consituent");

if (!(llvm::all_equal(constituents)))
if (!llvm::all_equal(constituents))
return rewriter.notifyMatchFailure(op, "composite is not splat");

auto splatConstituent = dyn_cast<FlatSymbolRefAttr>(constituents[0]);
Expand Down
Loading