Skip to content

[mlir][spirv] Extend capabilities and extensions requirements checking #68031

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

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[mlir][spirv] Extend capabilities and extensions requirements checking #68031

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
67 changes: 61 additions & 6 deletions mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,9 +43,13 @@ using namespace mlir;
template <typename LabelT>
static LogicalResult checkExtensionRequirements(
LabelT label, const spirv::TargetEnv &targetEnv,
const spirv::SPIRVType::ExtensionArrayRefVector &candidates) {
const spirv::SPIRVType::ExtensionArrayRefVector &candidates,
Comment on lines 44 to +46
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instead of adding this new parameter, could we pass in only a subset of candidates to this function? IE prune it on the caller side.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, that is also possible. Wanted to keep the signature of all 3 of the requirement checking function similar.
Also, wanted to have the prunning logic in one place.

Should I change the signature of this function?

ArrayRef<spirv::Extension> elidedCandidates = {}) {
for (const auto &ors : candidates) {
if (targetEnv.allows(ors))
if (targetEnv.allows(ors) ||
llvm::any_of(elidedCandidates, [&ors](spirv::Extension elidedExt) {
return llvm::is_contained(ors, elidedExt);
}))
continue;

LLVM_DEBUG({
Expand All @@ -71,9 +75,13 @@ static LogicalResult checkExtensionRequirements(
template <typename LabelT>
static LogicalResult checkCapabilityRequirements(
LabelT label, const spirv::TargetEnv &targetEnv,
const spirv::SPIRVType::CapabilityArrayRefVector &candidates) {
const spirv::SPIRVType::CapabilityArrayRefVector &candidates,
ArrayRef<spirv::Capability> elidedCandidates = {}) {
for (const auto &ors : candidates) {
if (targetEnv.allows(ors))
if (targetEnv.allows(ors) ||
llvm::any_of(elidedCandidates, [&ors](spirv::Capability elidedCap) {
return llvm::is_contained(ors, elidedCap);
}))
continue;

LLVM_DEBUG({
Expand All @@ -90,8 +98,55 @@ static LogicalResult checkCapabilityRequirements(
return success();
}

/// Returns true if the given `storageClass` needs explicit layout when used in
/// Shader environments.
/// Check capabilities and extensions requirements
/// Checks that `capCandidates`, `extCandidates`, and capability
/// (`capCandidates`) infered extension requirements are possible to be
/// satisfied with the given `targetEnv`.
/// It also provides a way to relax requirements for certain capabilities and
/// extensions (e.g., `elidedCapCandidates`, `elidedExtCandidates`), this is to
/// allow passes to relax certain requirements based on an option (e.g.,
/// relaxing bitwidth requirement, see `convertScalarType()`,
/// `ConvertVectorType()`).
template <typename LabelT>
static LogicalResult checkCapabilityAndExtensionRequirements(
LabelT label, const spirv::TargetEnv &targetEnv,
const spirv::SPIRVType::CapabilityArrayRefVector &capCandidates,
const spirv::SPIRVType::ExtensionArrayRefVector &extCandidates,
ArrayRef<spirv::Capability> elidedCapCandidates = {},
ArrayRef<spirv::Extension> elidedExtCandidates = {}) {
SmallVector<ArrayRef<spirv::Extension>, 8> updatedExtCandidates;
llvm::append_range(updatedExtCandidates, extCandidates);

if (failed(checkCapabilityRequirements(label, targetEnv, capCandidates,
elidedCapCandidates)))
return failure();
// Add capablity infered extensions to the list of extension requirement list,
// only considers the capabilities that already available in the `targetEnv`.

// WARNING: Some capabilities are part of both the core SPIR-V
// specification and an extension (e.g., 'Groups' capability is part of both
// core specification and SPV_AMD_shader_ballot extension, hence we should
// relax the capability inferred extension for these cases).
static const spirv::Capability multiModalCaps[] = {spirv::Capability::Groups};
ArrayRef<spirv::Capability> multiModalCapsArrayRef(multiModalCaps,
std::size(multiModalCaps));

for (auto cap : targetEnv.getAttr().getCapabilities()) {
if (llvm::any_of(multiModalCapsArrayRef,
[&cap](spirv::Capability mMCap) { return cap == mMCap; }))
continue;
std::optional<ArrayRef<spirv::Extension>> ext = getExtensions(cap);
if (ext)
updatedExtCandidates.push_back(*ext);
}
if (failed(checkExtensionRequirements(label, targetEnv, updatedExtCandidates,
elidedExtCandidates)))
return failure();
return success();
}

/// Returns true if the given `storageClass` needs explicit layout when used
/// in Shader environments.
static bool needsExplicitLayout(spirv::StorageClass storageClass) {
switch (storageClass) {
case spirv::StorageClass::PhysicalStorageBuffer:
Expand Down