/*
* Copyright 2024 Google LLC
*
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*/
#include "src/sksl/analysis/SkSLSpecialization.h"
#include "include/private/base/SkAssert.h"
#include "include/private/base/SkSpan_impl.h"
#include "src/sksl/SkSLAnalysis.h"
#include "src/sksl/SkSLDefines.h"
#include "src/sksl/analysis/SkSLProgramVisitor.h"
#include "src/sksl/ir/SkSLExpression.h"
#include "src/sksl/ir/SkSLFieldAccess.h"
#include "src/sksl/ir/SkSLFunctionCall.h"
#include "src/sksl/ir/SkSLFunctionDeclaration.h"
#include "src/sksl/ir/SkSLFunctionDefinition.h"
#include "src/sksl/ir/SkSLProgram.h"
#include "src/sksl/ir/SkSLProgramElement.h"
#include "src/sksl/ir/SkSLVariable.h"
#include "src/sksl/ir/SkSLVariableReference.h"
#include <algorithm>
#include <memory>
using namespace skia_private;
namespace SkSL::Analysis {
static bool parameter_mappings_are_equal(
const SpecializedParameters& left,
const SpecializedParameters& right) {
if (left.count() != right.count()) {
return false;
}
for (
const auto& [key, leftExpr] : left) {
const Expression** rightExpr = right.find(key);
if (!rightExpr) {
return false;
}
if (!Analysis::IsSameExpressionTree(*leftExpr, **rightExpr)) {
return false;
}
}
return true;
}
void FindFunctionsToSpecialize(
const Program& program,
SpecializationInfo* info,
const ParameterMatchesFn& parameterMatchesFn) {
class Searcher :
public ProgramVisitor {
public:
using ProgramVisitor::visitProgramElement;
using INHERITED = ProgramVisitor;
Searcher(SpecializationInfo& info,
const ParameterMatchesFn& parameterMatche
sFn)
: fSpecializationMap(info.fSpecializationMap)
, fSpecializedCallMap(info.fSpecializedCallMap)
, fParameterMatchesFn(parameterMatchesFn) {}
bool visitExpression(const Expression& expr) override {
if (expr.is<FunctionCall>()) {
const FunctionCall& call = expr.as<FunctionCall>();
const FunctionDeclaration& decl = call.function();
if (!decl.isIntrinsic()) {
SpecializedParameters specialization;
const int numParams = decl.parameters().size();
SkASSERT(call.arguments().size() == numParams);
for (int i = 0; i < numParams; i++) {
const Expression& arg = *call.arguments()[i];
// Specializations can only be made on arguments that are not complex
// expressions but only a variable reference or field access since these
// references will be inlined in the generated specialized functions.
const Variable* argBase = nullptr;
if (arg.is<VariableReference>()) {
argBase = arg.as<VariableReference>().variable();
} else if (arg.is<FieldAccess>() &&
arg.as<FieldAccess>().base()->is<VariableReference>()) {
argBase =
arg.as<FieldAccess>().base()->as<VariableReference>().variable();
} else {
continue;
}
SkASSERT(argBase);
const Variable* param = decl.parameters()[i];
// Check that this parameter fits the criteria to create a specialization.
if (!fParameterMatchesFn(*param)) {
continue;
}
if (argBase->storage() == Variable::Storage::kGlobal) {
specialization[param] = &arg;
} else if (argBase->storage() == Variable::Storage::kParameter) {
const Expression** uniformExpr =
fInheritedSpecializations.find(argBase);
SkASSERT(uniformExpr);
specialization[param] = *uniformExpr;
} else {
// TODO(b/353532475): Report an error instead of aborting.
SK_ABORT("Specialization requires a uniform or parameter variable");
}
}
// Only create a specialization for this function if there are
// variables to specialize on.
if (specialization.count() > 0) {
Specializations& specializations = fSpecializationMap[&decl];
SpecializedCallKey callKey{call.stablePointer(),
fInheritedSpecializationIndex};
for (int i = 0; i < specializations.size(); i++) {
const SpecializedParameters& entry = specializations[i];
if (parameter_mappings_are_equal(specialization, entry)) {
// This specialization has already been tracked.
fSpecializedCallMap[callKey] = i;
return INHERITED::visitExpression(expr);
}
}
// Set the index to the corresponding specialization this function call
// requires, also tracking the inherited specialization this function
// call is in so the right specialized function can be called.
SpecializationIndex specializationIndex = specializations.size();
fSpecializedCallMap[callKey] = specializationIndex;
specializations.push_back(specialization);
// We swap so we don't lose when our last inherited specializations were
// once we are done traversing this specific specialization.
fInheritedSpecializations.swap(specialization);
std::swap(fInheritedSpecializationIndex, specializationIndex);
this->visitProgramElement(*decl.definition());
std::swap(fInheritedSpecializationIndex, specializationIndex);
fInheritedSpecializations.swap(specialization);
} else {
// The function being called isn't specialized, but we need to walk the
// entire call graph or we may miss a specialized call entirely. Since
// nothing is specialized, it is safe to skip over repeated traversals.
if (!fVisitedFunctions.find(&decl)) {
fVisitedFunctions.add(&decl);
this->visitProgramElement(*decl.definition());
}
}
}
}
return INHERITED::visitExpression(expr);
}
private:
SpecializationMap& fSpecializationMap;
SpecializedCallMap& fSpecializedCallMap;
const ParameterMatchesFn& fParameterMatchesFn;
THashSet<const FunctionDeclaration*> fVisitedFunctions;
SpecializedParameters fInheritedSpecializations;
SpecializationIndex fInheritedSpecializationIndex = kUnspecialized;
};
for (const ProgramElement* elem : program.elements()) {
if (elem->is<FunctionDefinition>()) {
const FunctionDeclaration& decl = elem->as<FunctionDefinition>().declaration();
if (decl.isMain()) {
// Visit through the program call stack and aggregates any necessary
// function specializations.
Searcher(*info, parameterMatchesFn).visitProgramElement(*elem);
continue;
}
// Look for any function parameter which needs specialization.
for (const Variable* param : decl.parameters()) {
if (parameterMatchesFn(*param)) {
// We found a function that requires specialization. Ensure that this function
// ends up in the specialization map, whether or not it is reachable from
// main().
//
// Doing this here allows unreachable specialized functions to be discarded,
// because it will be in the specialization map with an array of zero necessary
// specializations to emit. If we didn't add this function to the specialization
// map at all, the code generator would try to emit it without applying
// specializations, and generally this would lead to invalid code.
info->fSpecializationMap[&decl];
break;
}
}
}
}
}
SpecializationIndex FindSpecializationIndexForCall(const FunctionCall& call,
const SpecializationInfo& info,
SpecializationIndex parentSpecializationIndex) {
SpecializedCallKey callKey{call.stablePointer(), parentSpecializationIndex};
SpecializationIndex* foundIndex = info.fSpecializedCallMap.find(callKey);
return foundIndex ? *foundIndex : kUnspecialized;
}
SkBitSet FindSpecializedParametersForFunction(const FunctionDeclaration& func,
const SpecializationInfo& info) {
SkBitSet result(func.parameters().size());
if (const Specializations* specializations = info.fSpecializationMap.find(&func)) {
const Analysis::SpecializedParameters& specializedParams = specializations->front();
const SkSpan<Variable* const> funcParams = func.parameters();
for (size_t index = 0; index < funcParams.size(); ++index) {
if (specializedParams.find(funcParams[index])) {
result.set(index);
}
}
}
return result;
}
void GetParameterMappingsForFunction(const FunctionDeclaration& func,
const SpecializationInfo& info,
SpecializationIndex specializationIndex,
const ParameterMappingCallback& callback) {
if (specializationIndex != Analysis::kUnspecialized) {
if (const Specializations* specializations = info.fSpecializationMap.find(&func)) {
const Analysis::SpecializedParameters& specializedParams =
specializations->at(specializationIndex);
const SkSpan<Variable* const> funcParams = func.parameters();
for (size_t index = 0; index < funcParams.size(); ++index) {
const Variable* param = funcParams[index];
if (const Expression** expr = specializedParams.find(param)) {
callback(index, param, *expr);
}
}
}
}
}
} // namespace SkSL::Analysis
