Skip to content

Commit

Permalink
Rewrite some handling codes of unpacked output locations
Browse files Browse the repository at this point in the history
We have the logic to keep all locations in some conditions to initialize
more location map entries than actually needed. The original logic is
somewhat vague when TCS dynamic indexing is involved (the same to mesh
shader outputs). For example:

  layout(location = 10) out float data[5]
  ...
  data[i] = XXX

Location 0-9 should be initialized to 'unmapped' status and we only map
the actual location range, which is location 10-14. The original code
mapped all locations 0-14. This is incorrect if we have two arrays:

  layout(location = 5) out float data1[5]
  layout(location = 10) out float data2[5]
  ...
  data1[i] = XXX
  data2[j] = YYY

In resource collecting pass, also re-write some handling codes of
unpacked output locations. This is to address such case:

  layout(location = 0) out float data1[5]
  layout(location = 5) out float data2
  layout(location = 6) out float data3[5]
  ...
  data1[i] = XXX
  data2 = 0.5
  data3[j] = YYY

When handing dynamic indexing, data1 is assigned to location 0-5 and data3
is assigned to location 6-10 without location mapping needed in resource
collecting pass. data2 could take location 5 but our current logic assigns
it to 0 without checking already-occupied locations. Such error is found
when compiling a separate TCS (not full pipeline). So we must first
colllect occupied locations and do location mapping after that.
  • Loading branch information
amdrexu committed Dec 22, 2023
1 parent 217a8e3 commit f5542ee
Show file tree
Hide file tree
Showing 2 changed files with 152 additions and 48 deletions.
94 changes: 71 additions & 23 deletions lgc/builder/InOutBuilder.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -181,7 +181,7 @@ Value *BuilderImpl::readGenericInputOutput(bool isOutput, Type *resultTy, unsign
assert(isOutput == false || m_shaderStage == ShaderStageTessControl);

// Fold constant locationOffset into location. (Currently a variable locationOffset is only supported in
// TCS, TES, and FS custom interpolation.)
// TCS, TES, mesh shader, and FS custom interpolation.)
bool isDynLocOffset = true;
if (auto constLocOffset = dyn_cast<ConstantInt>(locationOffset)) {
location += constLocOffset->getZExtValue();
Expand Down Expand Up @@ -267,8 +267,8 @@ Instruction *BuilderImpl::CreateWriteGenericOutput(Value *valueToWrite, unsigned
Value *vertexOrPrimitiveIndex) {
assert(valueToWrite->getType()->isAggregateType() == false);

// Fold constant locationOffset into location. (Currently a variable locationOffset is only supported in
// TCS.)
// Fold constant locationOffset into location (Currently a variable locationOffset is only supported in
// TCS or mesh shader).
bool isDynLocOffset = true;
if (auto constLocOffset = dyn_cast<ConstantInt>(locationOffset)) {
location += constLocOffset->getZExtValue();
Expand Down Expand Up @@ -346,7 +346,7 @@ Instruction *BuilderImpl::CreateWriteGenericOutput(Value *valueToWrite, unsigned
//
// @param isOutput : False for input, true for output
// @param location : Input/output base location
// @param locationCount : Count of locations taken by the input
// @param locationCount : Count of locations taken by the input/output
// @param inOutInfo : Extra input/output information
// @param vertexOrPrimIndex : For TCS/TES/GS/mesh shader per-vertex input/output: vertex index;
// for mesh shader per-primitive output: primitive index;
Expand Down Expand Up @@ -389,44 +389,92 @@ void BuilderImpl::markGenericInputOutputUsage(bool isOutput, unsigned location,
}
}

if (!isOutput || m_shaderStage != ShaderStageGeometry) {
if (!(m_shaderStage == ShaderStageGeometry && isOutput)) {
// Not GS output
bool keepAllLocations = false;
if (getPipelineState()->isUnlinked()) {
if (isOutput && m_shaderStage != ShaderStageFragment) {
ShaderStage nextStage = m_pipelineState->getNextShaderStage(m_shaderStage);
keepAllLocations = nextStage == ShaderStageFragment || nextStage == ShaderStageInvalid;
if (isOutput) {
// Keep all locations if the next stage of the output is fragment shader or is unspecified
if (m_shaderStage != ShaderStageFragment) {
ShaderStage nextStage = m_pipelineState->getNextShaderStage(m_shaderStage);
keepAllLocations = nextStage == ShaderStageFragment || nextStage == ShaderStageInvalid;
}
} else {
// Keep all locations if it is the input of fragment shader
keepAllLocations = m_shaderStage == ShaderStageFragment;
}
if (m_shaderStage == ShaderStageFragment && !isOutput)
keepAllLocations = true;
}
unsigned startLocation = (keepAllLocations ? 0 : location);
// NOTE: The non-invalid value as initial new Location info or new location is used to identify the dynamic indexing
// location.
// Non-GS-output case.

if (inOutLocInfoMap) {
for (unsigned i = startLocation; i < location + locationCount; ++i) {
// Handle per-vertex input/output
if (keepAllLocations) {
// If keeping all locations, add location map entries whose locations are before this input/output
for (unsigned i = 0; i < location; ++i) {
InOutLocationInfo origLocationInfo;
origLocationInfo.setLocation(i);
if (inOutLocInfoMap->count(origLocationInfo) == 0) {
// Add this location map entry only if it doesn't exist
auto &newLocationInfo = (*inOutLocInfoMap)[origLocationInfo];
newLocationInfo.setData(InvalidValue);
}
}
}

// Add location map entries for this input/output
for (unsigned i = 0; i < locationCount; ++i) {
InOutLocationInfo origLocationInfo;
origLocationInfo.setLocation(i);
origLocationInfo.setLocation(location + i);
origLocationInfo.setComponent(inOutInfo.getComponent());
auto &newLocationInfo = (*inOutLocInfoMap)[origLocationInfo];
newLocationInfo.setData(isDynLocOffset ? i : InvalidValue);
if (isDynLocOffset) {
// When dynamic indexing, map the location directly
newLocationInfo.setLocation(location + i);
newLocationInfo.setComponent(inOutInfo.getComponent());
} else
newLocationInfo.setData(InvalidValue);
}
}

if (perPatchInOutLocMap) {
for (unsigned i = startLocation; i < location + locationCount; ++i)
(*perPatchInOutLocMap)[i] = isDynLocOffset ? i : InvalidValue;
// Handle per-patch input/output
if (keepAllLocations) {
// If keeping all locations, add location map entries whose locations are before this input/output
for (unsigned i = 0; i < location; ++i) {
// Add this location map entry only if it doesn't exist
if (perPatchInOutLocMap->count(i) == 0)
(*perPatchInOutLocMap)[i] = InvalidValue;
}
}

// Add location map entries for this input/output
for (unsigned i = 0; i < locationCount; ++i)
(*perPatchInOutLocMap)[location + i] =
isDynLocOffset ? location + i : InvalidValue; // When dynamic indexing, map the location
}

if (perPrimitiveInOutLocMap) {
for (unsigned i = startLocation; i < location + locationCount; ++i)
(*perPrimitiveInOutLocMap)[i] = isDynLocOffset ? i : InvalidValue;
// Handle per-primitive input/output
if (keepAllLocations) {
// If keeping all locations, add location map entries whose locations are before this input/output
for (unsigned i = 0; i < location; ++i) {
// Add this location map entry only if it doesn't exist
if (perPrimitiveInOutLocMap->count(i) == 0)
(*perPrimitiveInOutLocMap)[i] = InvalidValue;
}
}

// Add location map entries for this input/output
for (unsigned i = 0; i < locationCount; ++i)
(*perPrimitiveInOutLocMap)[location + i] =
isDynLocOffset ? location + i : InvalidValue; // When dynamic indexing, map the location
}
} else {
// GS output. We include the stream ID with the location in the map key.
// GS output
for (unsigned i = 0; i < locationCount; ++i) {
InOutLocationInfo outLocationInfo;
outLocationInfo.setLocation(location + i);
outLocationInfo.setComponent(inOutInfo.getComponent());
outLocationInfo.setStreamId(inOutInfo.getStreamId());
outLocationInfo.setStreamId(inOutInfo.getStreamId()); // Include the stream ID in the map key.
auto &newLocationInfo = (*inOutLocInfoMap)[outLocationInfo];
newLocationInfo.setData(InvalidValue);
}
Expand Down
106 changes: 81 additions & 25 deletions lgc/patch/PatchResourceCollect.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -2906,68 +2906,124 @@ void PatchResourceCollect::updateOutputLocInfoMapWithUnpack() {

// Update the value of outputLocInfoMap
if (!outputLocInfoMap.empty()) {
unsigned nextMapLoc[MaxGsStreams] = {};
DenseMap<unsigned, unsigned> alreadyMappedLocs[MaxGsStreams]; // Map from original location to new location
unsigned nextAvailableLoc[MaxGsStreams] = {};
DenseMap<unsigned, unsigned> locMap[MaxGsStreams]; // Map from original location to new location
DenseSet<unsigned> occupiedLocs[MaxGsStreams]; // Collection of already-occupied locations in location mapping

// Collect all mapped locations before we do location mapping for those unmapped
for (auto &locInfoPair : outputLocInfoMap) {
const auto &locationInfo = locInfoPair.first;
auto &newLocationInfo = locInfoPair.second;

const unsigned streamId = m_shaderStage == ShaderStageGeometry ? locationInfo.getStreamId() : 0;

if (!newLocationInfo.isInvalid()) {
// Record mapped locations
const unsigned locAlreadyMapped = locationInfo.getLocation();
const unsigned newLocMappedTo = newLocationInfo.getLocation();
assert(newLocMappedTo != InvalidValue);

locMap[streamId][locAlreadyMapped] = newLocMappedTo;
occupiedLocs[streamId].insert(newLocMappedTo);
}
}

// Do location mapping
for (auto &locInfoPair : outputLocInfoMap) {
const auto &locationInfo = locInfoPair.first;
auto &newLocationInfo = locInfoPair.second;

if (!newLocationInfo.isInvalid())
continue; // Skip any location that is mapped
continue; // Skip mapped locations

const unsigned streamId = m_shaderStage == ShaderStageGeometry ? locationInfo.getStreamId() : 0;

newLocationInfo.setData(0);
newLocationInfo.setComponent(locationInfo.getComponent());
const unsigned streamId = m_shaderStage == ShaderStageGeometry ? locationInfo.getStreamId() : 0;
newLocationInfo.setStreamId(streamId);

const unsigned locToBeMapped = locationInfo.getLocation();
unsigned mappedLoc = InvalidValue;
unsigned newLocMappedTo = InvalidValue;
const bool keepLocation = m_shaderStage == ShaderStageGeometry && !canChangeOutputLocationsForGs();
if (keepLocation) {
// Keep location unchanged
mappedLoc = locToBeMapped;
newLocMappedTo = locToBeMapped;
} else {
// Map to new location
if (alreadyMappedLocs[streamId].count(locToBeMapped) > 0) {
mappedLoc = alreadyMappedLocs[streamId][locToBeMapped];
if (locMap[streamId].count(locToBeMapped) > 0) {
newLocMappedTo = locMap[streamId][locToBeMapped];
} else {
mappedLoc = nextMapLoc[streamId]++;
do {
// Try to find a new location that has not been occupied
newLocMappedTo = nextAvailableLoc[streamId]++;
} while (occupiedLocs[streamId].count(newLocMappedTo) > 0);

// NOTE: Record the map because we are handling multiple pairs of <location, component>. Some pairs have the
// same location while the components are different.
alreadyMappedLocs[streamId][locToBeMapped] = mappedLoc;
locMap[streamId][locToBeMapped] = newLocMappedTo;
occupiedLocs[streamId].insert(newLocMappedTo);
}
}

assert(mappedLoc != InvalidValue);
newLocationInfo.setLocation(mappedLoc);
assert(newLocMappedTo != InvalidValue);
newLocationInfo.setLocation(newLocMappedTo);

if (m_shaderStage == ShaderStageGeometry)
inOutUsage.gs.outLocCount[streamId] = std::max(inOutUsage.gs.outLocCount[streamId], mappedLoc + 1);
inOutUsage.gs.outLocCount[streamId] = std::max(inOutUsage.gs.outLocCount[streamId], newLocMappedTo + 1);
}
}

// Update the value of perPatchOutputLocMap
if (!perPatchOutputLocMap.empty()) {
unsigned nextMapLoc = 0;
assert(m_shaderStage == ShaderStageTessControl);

unsigned nextAvailableLoc = 0;
DenseSet<unsigned> occupiedLocs; // Collection of already-occupied locations in location mapping

// Collect all mapped locations before we do location mapping for those unmapped
for (auto &locPair : perPatchOutputLocMap) {
if (locPair.second != InvalidValue)
occupiedLocs.insert(locPair.second); // Record mapped locations
}

// Do location mapping
for (auto &locPair : perPatchOutputLocMap) {
if (locPair.second == InvalidValue) {
// Only do location mapping if the per-patch output has not been mapped
locPair.second = nextMapLoc++;
} else
assert(m_shaderStage == ShaderStageTessControl);
if (locPair.second != InvalidValue)
continue; // Skip mapped locations

unsigned newLocMappedTo = InvalidValue;
do {
// Try to find a new location that has not been occupied
newLocMappedTo = nextAvailableLoc++;
} while (occupiedLocs.count(newLocMappedTo) > 0);
locPair.second = newLocMappedTo;
}
}

// Update the value of perPrimitiveOutputLocMap
if (!perPrimitiveOutputLocMap.empty()) {
unsigned nextMapLoc = 0;
assert(m_shaderStage == ShaderStageMesh);

unsigned nextAvailableLoc = 0;
DenseSet<unsigned> occupiedLocs; // Collection of already-occupied locations in location mapping

// Collect all mapped locations before we do location mapping for those unmapped
for (auto &locPair : perPrimitiveOutputLocMap) {
if (locPair.second != InvalidValue)
occupiedLocs.insert(locPair.second); // Record mapped locations
}

// Do location mapping
for (auto &locPair : perPrimitiveOutputLocMap) {
if (locPair.second == InvalidValue) {
// Only do location mapping if the per-primitive output has not been mapped
locPair.second = nextMapLoc++;
} else
assert(m_shaderStage == ShaderStageMesh);
if (locPair.second != InvalidValue)
continue; // Skip mapped locations

unsigned newLocMappedTo = InvalidValue;
do {
// Try to find a new location that has not been occupied
newLocMappedTo = nextAvailableLoc++;
} while (occupiedLocs.count(newLocMappedTo) > 0);
locPair.second = newLocMappedTo;
}
}

Expand Down

0 comments on commit f5542ee

Please sign in to comment.