/*
* Copyright 2018 Google Inc.
*
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*/
#ifndef SKSL_JIT
#define SKSL_JIT
#ifdef SK_LLVM_AVAILABLE
#include "ir/SkSLBinaryExpression.h"
#include "ir/SkSLBreakStatement.h"
#include "ir/SkSLContinueStatement.h"
#include "ir/SkSLExpression.h"
#include "ir/SkSLDoStatement.h"
#include "ir/SkSLForStatement.h"
#include "ir/SkSLFunctionCall.h"
#include "ir/SkSLFunctionDefinition.h"
#include "ir/SkSLIfStatement.h"
#include "ir/SkSLIndexExpression.h"
#include "ir/SkSLPrefixExpression.h"
#include "ir/SkSLPostfixExpression.h"
#include "ir/SkSLProgram.h"
#include "ir/SkSLReturnStatement.h"
#include "ir/SkSLStatement.h"
#include "ir/SkSLSwizzle.h"
#include "ir/SkSLTernaryExpression.h"
#include "ir/SkSLVarDeclarationsStatement.h"
#include "ir/SkSLVariableReference.h"
#include "ir/SkSLWhileStatement.h"
#include "llvm-c/Analysis.h"
#include "llvm-c/Core.h"
#include "llvm-c/OrcBindings.h"
#include "llvm-c/Support.h"
#include "llvm-c/Target.h"
#include "llvm-c/Transforms/PassManagerBuilder.h"
#include "llvm-c/Types.h"
#include <stack>
class SkRasterPipeline;
namespace SkSL {
struct AppendStage;
/**
* A just-in-time compiler for SkSL code which uses an LLVM backend. Only available when the
* skia_llvm_path gn arg is set.
*
* Example of using SkSLJIT to set up an SkJumper pipeline stage:
*
* #ifdef SK_LLVM_AVAILABLE
* SkSL::Compiler compiler;
* SkSL::Program::Settings settings;
* std::unique_ptr<SkSL::Program> program = compiler.convertProgram(
SkSL::Program::kPipelineStage_Kind,
* "void swap(int x, int y, inout float4 color) {"
* " color.rb = color.br;"
* "}",
* settings);
* if (!program) {
* printf("%s\n", compiler.errorText().c_str());
* abort();
* }
* SkSL::JIT& jit = *scratch->make<SkSL::JIT>(&compiler);
* std::unique_ptr<SkSL::JIT::Module> module = jit.compile(std::move(program));
* void* func = module->getJumperStage("swap");
* p->append(func, nullptr);
* #endif
*/
class JIT {
typedef int StackIndex;
public:
class Module {
public:
/**
* Returns the address of a symbol in the module.
*/
void* getSymbol(const char* name);
/**
* Returns the address of a function as an SkJumper pipeline stage. The function must have
* the signature void <name>(int x, int y, inout float4 color). The returned function will
* have the correct signature to function as an SkJumper stage (meaning it will actually
* have a different signature at runtime, accepting vector parameters and operating on
* multiple pixels simultaneously as is normal for SkJumper stages).
*/
void* getJumperStage(const char* name);
~Module() {
LLVMOrcDisposeSharedModuleRef(fSharedModule);
}
private:
Module(std::unique_ptr<Program> program,
LLVMSharedModuleRef sharedModule,
LLVMOrcJITStackRef jitStack)
: fProgram(std::move(program))
, fSharedModule(sharedModule)
, fJITStack(jitStack) {}
std::unique_ptr<Program> fProgram;
LLVMSharedModuleRef fSharedModule;
LLVMOrcJITStackRef fJITStack;
friend class JIT;
};
JIT(Compiler* compiler);
~JIT();
/**
* Just-in-time compiles an SkSL program and returns the resulting Module. The JIT must not be
* destroyed before all of its Modules are destroyed.
*/
std::unique_ptr<Module> compile(std::unique_ptr<Program> program);
private:
static constexpr int CHANNELS = 4;
enum TypeKind {
kFloat_TypeKind,
kInt_TypeKind,
kUInt_TypeKind,
kBool_TypeKind
};
class LValue {
public:
virtual ~LValue() {}
virtual LLVMValueRef load(LLVMBuilderRef builder) = 0;
virtual void store(LLVMBuilderRef builder, LLVMValueRef value) = 0;
};
void addBuiltinFunction(const char* ourName, const char* realName, LLVMTypeRef returnType,
std::vector<LLVMTypeRef> parameters);
void loadBuiltinFunctions();
void setBlock(LLVMBuilderRef builder, LLVMBasicBlockRef block);
LLVMTypeRef getType(const Type& type);
TypeKind typeKind(const Type& type);
std::unique_ptr<LValue> getLValue(LLVMBuilderRef builder, const Expression& expr);
void vectorize(LLVMBuilderRef builder, LLVMValueRef* value, int columns);
void vectorize(LLVMBuilderRef builder, const BinaryExpression& b, LLVMValueRef* left,
LLVMValueRef* right);
LLVMValueRef compileBinary(LLVMBuilderRef builder, const BinaryExpression& b);
LLVMValueRef compileConstructor(LLVMBuilderRef builder, const Constructor& c);
LLVMValueRef compileFunctionCall(LLVMBuilderRef builder, const FunctionCall& fc);
LLVMValueRef compileIndex(LLVMBuilderRef builder, const IndexExpression& v);
LLVMValueRef compilePostfix(LLVMBuilderRef builder, const PostfixExpression& p);
LLVMValueRef compilePrefix(LLVMBuilderRef builder, const PrefixExpression& p);
LLVMValueRef compileSwizzle(LLVMBuilderRef builder, const Swizzle& s);
LLVMValueRef compileVariableReference(LLVMBuilderRef builder, const VariableReference& v);
LLVMValueRef compileTernary(LLVMBuilderRef builder, const TernaryExpression& t);
LLVMValueRef compileExpression(LLVMBuilderRef builder, const Expression& expr);
void appendStage(LLVMBuilderRef builder, const AppendStage& a);
void compileBlock(LLVMBuilderRef builder, const Block& block);
void compileBreak(LLVMBuilderRef builder, const BreakStatement& b);
void compileContinue(LLVMBuilderRef builder, const ContinueStatement& c);
void compileDo(LLVMBuilderRef builder, const DoStatement& d);
void compileFor(LLVMBuilderRef builder, const ForStatement& f);
void compileIf(LLVMBuilderRef builder, const IfStatement& i);
void compileReturn(LLVMBuilderRef builder, const ReturnStatement& r);
void compileVarDeclarations(LLVMBuilderRef builder, const VarDeclarationsStatement& decls);
void compileWhile(LLVMBuilderRef builder, const WhileStatement& w);
void compileStatement(LLVMBuilderRef builder, const Statement& stmt);
// The "Vector" variants of functions attempt to compile a given expression or statement as part
// of a vectorized SkJumper stage function - that is, with r, g, b, and a each being vectors of
// fVectorCount floats. So a statement like "color.r = 0;" looks like it modifies a single
// channel of a single pixel, but the compiled code will actually modify the red channel of
// fVectorCount pixels at once.
//
// As not everything can be vectorized, these calls return a bool to indicate whether they were
// successful. If anything anywhere in the function cannot be vectorized, the JIT will fall back
// to looping over the pixels instead.
//
// Since we process multiple pixels at once, and each pixel consists of multiple color channels,
// expressions may effectively result in a vector-of-vectors. We produce zero to four outputs
// when compiling expression, each of which is a vector, so that e.g. float2(1, 0) actually
// produces two vectors, one containing all 1s, the other all 0s. The out parameter always
// allows for 4 channels, but the functions produce 0 to 4 channels depending on the type they
// are operating on. Thus evaluating "color.rgb" actually fills in out[0] through out[2],
// leaving out[3] uninitialized.
// As the number of outputs can be inferred from the type of the expression, it is not
// explicitly signalled anywhere.
bool compileVectorBinary(LLVMBuilderRef builder, const BinaryExpression& b,
LLVMValueRef out[CHANNELS]);
bool compileVectorConstructor(LLVMBuilderRef builder, const Constructor& c,
LLVMValueRef out[CHANNELS]);
bool compileVectorFloatLiteral(LLVMBuilderRef builder, const FloatLiteral& f,
LLVMValueRef out[CHANNELS]);
bool compileVectorSwizzle(LLVMBuilderRef builder, const Swizzle& s,
LLVMValueRef out[CHANNELS]);
bool compileVectorVariableReference(LLVMBuilderRef builder, const VariableReference& v,
LLVMValueRef out[CHANNELS]);
bool compileVectorExpression(LLVMBuilderRef builder, const Expression& expr,
LLVMValueRef out[CHANNELS]);
bool getVectorLValue(LLVMBuilderRef builder, const Expression& e, LLVMValueRef out[CHANNELS]);
/**
* Evaluates the left and right operands of a binary operation, promoting one of them to a
* vector if necessary to make the types match.
*/
bool getVectorBinaryOperands(LLVMBuilderRef builder, const Expression& left,
LLVMValueRef outLeft[CHANNELS], const Expression& right,
LLVMValueRef outRight[CHANNELS]);
bool compileVectorStatement(LLVMBuilderRef builder, const Statement& stmt);
/**
* Returns true if this function has the signature void(int, int, inout float4) and thus can be
* used as an SkJumper stage.
*/
bool hasStageSignature(const FunctionDeclaration& f);
/**
* Attempts to compile a vectorized stage function, returning true on success. A stage function
* of e.g. "color.r = 0;" will produce code which sets the entire red vector to zeros in a
* single instruction, thus calculating several pixels at once.
*/
bool compileStageFunctionVector(const FunctionDefinition& f, LLVMValueRef newFunc);
/**
* Fallback function which loops over the pixels, for when vectorization fails. A stage function
* of e.g. "color.r = 0;" will produce a loop which iterates over the entries in the red vector,
* setting each one to zero individually.
*/
void compileStageFunctionLoop(const FunctionDefinition& f, LLVMValueRef newFunc);
/**
* Called when compiling a function which has the signature of an SkJumper stage. Produces a
* version of the function which can be plugged into SkJumper (thus having a signature which
* accepts four vectors, one for each color channel, containing the color data of multiple
* pixels at once). To go from SkSL code which operates on a single pixel at a time to CPU code
* which operates on multiple pixels at once, the code is either vectorized using
* compileStageFunctionVector or wrapped in a loop using compileStageFunctionLoop.
*/
LLVMValueRef compileStageFunction(const FunctionDefinition& f);
/**
* Compiles an SkSL function to an LLVM function. If the function has the signature of an
* SkJumper stage, it will *also* be compiled by compileStageFunction, resulting in both a stage
* and non-stage version of the function.
*/
LLVMValueRef compileFunction(const FunctionDefinition& f);
void createModule();
void optimize();
bool isColorRef(const Expression& expr);
static uint64_t resolveSymbol(const char* name, JIT* jit);
const char* fCPU;
int fVectorCount;
Compiler& fCompiler;
std::unique_ptr<Program> fProgram;
LLVMContextRef fContext;
LLVMModuleRef fModule;
LLVMSharedModuleRef fSharedModule;
LLVMOrcJITStackRef fJITStack;
LLVMValueRef fCurrentFunction;
LLVMBasicBlockRef fAllocaBlock;
LLVMBasicBlockRef fCurrentBlock;
LLVMTypeRef fVoidType;
LLVMTypeRef fInt1Type;
LLVMTypeRef fInt1VectorType;
LLVMTypeRef fInt1Vector2Type;
LLVMTypeRef fInt1Vector3Type;
LLVMTypeRef fInt1Vector4Type;
LLVMTypeRef fInt8Type;
LLVMTypeRef fInt8PtrType;
LLVMTypeRef fInt32Type;
LLVMTypeRef fInt32VectorType;
LLVMTypeRef fInt32Vector2Type;
LLVMTypeRef fInt32Vector3Type;
LLVMTypeRef fInt32Vector4Type;
LLVMTypeRef fInt64Type;
LLVMTypeRef fSizeTType;
LLVMTypeRef fFloat32Type;
LLVMTypeRef fFloat32VectorType;
LLVMTypeRef fFloat32Vector2Type;
LLVMTypeRef fFloat32Vector3Type;
LLVMTypeRef fFloat32Vector4Type;
// Our SkSL stage functions have a single float4 for color, but the actual SkJumper stage
// function has four separate vectors, one for each channel. These four values are references to
// the red, green, blue, and alpha vectors respectively.
LLVMValueRef fChannels[CHANNELS];
// when processing a stage function, this points to the SkSL color parameter (an inout float4)
const Variable* fColorParam;
std::unordered_map<const FunctionDeclaration*, LLVMValueRef> fFunctions;
std::unordered_map<const Variable*, LLVMValueRef> fVariables;
// LLVM function parameters are read-only, so when modifying function parameters we need to
// first promote them to variables. This keeps track of which parameters have been promoted.
std::set<const Variable*> fPromotedParameters;
std::vector<LLVMBasicBlockRef> fBreakTarget;
std::vector<LLVMBasicBlockRef> fContinueTarget;
LLVMValueRef fFoldAnd2Func;
LLVMValueRef fFoldOr2Func;
LLVMValueRef fFoldAnd3Func;
LLVMValueRef fFoldOr3Func;
LLVMValueRef fFoldAnd4Func;
LLVMValueRef fFoldOr4Func;
LLVMValueRef fAppendFunc;
LLVMValueRef fAppendCallbackFunc;
LLVMValueRef fDebugFunc;
};
} // namespace
#endif // SK_LLVM_AVAILABLE
#endif // SKSL_JIT