Skip to content

LLama inference implementation in Scala #82

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
szymon-rd wants to merge 1 commit into
base: main
Choose a base branch
from
+9,066 −137
Open

LLama inference implementation in Scala #82

Show file tree
Hide file tree
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
4 changes: 3 additions & 1 deletion .gitignore
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,4 +31,6 @@ metals.sbt
smoke
julia


# Model weight files
*.gguf
.lwjgl/
7 changes: 6 additions & 1 deletion build.sbt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -153,10 +153,15 @@ lazy val e2eTest = (project in file("cyfra-e2e-test"))
.settings(publish / skip := true)
.dependsOn(runtime, fs2interop, foton)

lazy val llama = (project in file("cyfra-llama"))
.settings(commonSettings, runnerSettings)
.settings(publish / skip := true)
.dependsOn(runtime, dsl, core, utility)

lazy val root = (project in file("."))
.settings(name := "Cyfra")
.settings(publish / skip := true)
.aggregate(compiler, dsl, foton, core, runtime, vulkan, examples, fs2interop, fluids, analytics, utility, spirvTools, vscode)
.aggregate(compiler, dsl, foton, core, runtime, vulkan, examples, fs2interop, fluids, analytics, utility, spirvTools, vscode, llama)

e2eTest / Test / javaOptions ++= Seq("-Dorg.lwjgl.system.stackSize=1024", "-DuniqueLibraryNames=true")

Expand Down
11 changes: 10 additions & 1 deletion cyfra-compiler/src/main/scala/io/computenode/cyfra/spirv/Context.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ import io.computenode.cyfra.dsl.binding.{GBuffer, GUniform}
import io.computenode.cyfra.dsl.macros.FnCall.FnIdentifier
import io.computenode.cyfra.spirv.SpirvConstants.HEADER_REFS_TOP
import io.computenode.cyfra.spirv.compilers.FunctionCompiler.SprivFunction
import io.computenode.cyfra.spirv.compilers.SpirvProgramCompiler.ArrayBufferBlock
import io.computenode.cyfra.spirv.compilers.SpirvProgramCompiler.{ArrayBufferBlock, SharedBlock}
import izumi.reflect.Tag
import izumi.reflect.macrortti.LightTypeTag

Expand All @@ -13,15 +13,24 @@ private[cyfra] case class Context(
funPointerTypeMap: Map[Int, Int] = Map(),
uniformPointerMap: Map[Int, Int] = Map(),
inputPointerMap: Map[Int, Int] = Map(),
workgroupPointerMap: Map[Int, Int] = Map(),
funcTypeMap: Map[(LightTypeTag, List[LightTypeTag]), Int] = Map(),
voidTypeRef: Int = -1,
voidFuncTypeRef: Int = -1,
workerIndexRef: Int = -1,
localInvocationIndexRef: Int = -1,
localInvocationIdRef: Int = -1,
workgroupIdRef: Int = -1,
numWorkgroupsRef: Int = -1,
subgroupIdRef: Int = -1,
subgroupLocalInvocationIdRef: Int = -1,
subgroupSizeRef: Int = -1,
uniformVarRefs: Map[GUniform[?], Int] = Map.empty,
bindingToStructType: Map[Int, Int] = Map.empty,
constRefs: Map[(Tag[?], Any), Int] = Map(),
exprRefs: Map[Int, Int] = Map(),
bufferBlocks: Map[GBuffer[?], ArrayBufferBlock] = Map(),
sharedVarRefs: Map[Int, SharedBlock] = Map(),
nextResultId: Int = HEADER_REFS_TOP,
nextBinding: Int = 0,
exprNames: Map[Int, String] = Map(),
Expand Down
58 changes: 58 additions & 0 deletions cyfra-compiler/src/main/scala/io/computenode/cyfra/spirv/Opcodes.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -516,6 +516,20 @@ private[cyfra] object Opcodes:
val Reduce = Code("Reduce", 0)
val InclusiveScan = Code("InclusiveScan", 1)
val ExclusiveScan = Code("ExclusiveScan", 2)
val ClusteredReduce = Code("ClusteredReduce", 3)

object MemorySemantics:
val None = Code("None", 0x0)
val Acquire = Code("Acquire", 0x2)
val Release = Code("Release", 0x4)
val AcquireRelease = Code("AcquireRelease", 0x8)
val SequentiallyConsistent = Code("SequentiallyConsistent", 0x10)
val UniformMemory = Code("UniformMemory", 0x40)
val SubgroupMemory = Code("SubgroupMemory", 0x80)
val WorkgroupMemory = Code("WorkgroupMemory", 0x100)
val CrossWorkgroupMemory = Code("CrossWorkgroupMemory", 0x200)
val AtomicCounterMemory = Code("AtomicCounterMemory", 0x400)
val ImageMemory = Code("ImageMemory", 0x800)

object KernelEnqueueFlags:
val NoWait = Code("NoWait", 0)
Expand Down Expand Up @@ -589,6 +603,14 @@ private[cyfra] object Opcodes:
val SubgroupDispatch = Code("SubgroupDispatch", 58)
val NamedBarrier = Code("NamedBarrier", 59)
val PipeStorage = Code("PipeStorage", 60)
val GroupNonUniform = Code("GroupNonUniform", 61)
val GroupNonUniformVote = Code("GroupNonUniformVote", 62)
val GroupNonUniformArithmetic = Code("GroupNonUniformArithmetic", 63)
val GroupNonUniformBallot = Code("GroupNonUniformBallot", 64)
val GroupNonUniformShuffle = Code("GroupNonUniformShuffle", 65)
val GroupNonUniformShuffleRelative = Code("GroupNonUniformShuffleRelative", 66)
val GroupNonUniformClustered = Code("GroupNonUniformClustered", 67)
val GroupNonUniformQuad = Code("GroupNonUniformQuad", 68)
val SubgroupBallotKHR = Code("SubgroupBallotKHR", 4423)
val DrawParameters = Code("DrawParameters", 4427)
val SubgroupVoteKHR = Code("SubgroupVoteKHR", 4431)
Expand Down Expand Up @@ -949,6 +971,42 @@ private[cyfra] object Opcodes:
val OpSubgroupImageBlockReadINTEL = Code("OpSubgroupImageBlockReadINTEL", 5577)
val OpSubgroupImageBlockWriteINTEL = Code("OpSubgroupImageBlockWriteINTEL", 5578)

// GroupNonUniform operations (Vulkan 1.1+)
val OpGroupNonUniformElect = Code("OpGroupNonUniformElect", 333)
val OpGroupNonUniformAll = Code("OpGroupNonUniformAll", 334)
val OpGroupNonUniformAny = Code("OpGroupNonUniformAny", 335)
val OpGroupNonUniformAllEqual = Code("OpGroupNonUniformAllEqual", 336)
val OpGroupNonUniformBroadcast = Code("OpGroupNonUniformBroadcast", 337)
val OpGroupNonUniformBroadcastFirst = Code("OpGroupNonUniformBroadcastFirst", 338)
val OpGroupNonUniformBallot = Code("OpGroupNonUniformBallot", 339)
val OpGroupNonUniformInverseBallot = Code("OpGroupNonUniformInverseBallot", 340)
val OpGroupNonUniformBallotBitExtract = Code("OpGroupNonUniformBallotBitExtract", 341)
val OpGroupNonUniformBallotBitCount = Code("OpGroupNonUniformBallotBitCount", 342)
val OpGroupNonUniformBallotFindLSB = Code("OpGroupNonUniformBallotFindLSB", 343)
val OpGroupNonUniformBallotFindMSB = Code("OpGroupNonUniformBallotFindMSB", 344)
val OpGroupNonUniformShuffle = Code("OpGroupNonUniformShuffle", 345)
val OpGroupNonUniformShuffleXor = Code("OpGroupNonUniformShuffleXor", 346)
val OpGroupNonUniformShuffleUp = Code("OpGroupNonUniformShuffleUp", 347)
val OpGroupNonUniformShuffleDown = Code("OpGroupNonUniformShuffleDown", 348)
val OpGroupNonUniformIAdd = Code("OpGroupNonUniformIAdd", 349)
val OpGroupNonUniformFAdd = Code("OpGroupNonUniformFAdd", 350)
val OpGroupNonUniformIMul = Code("OpGroupNonUniformIMul", 351)
val OpGroupNonUniformFMul = Code("OpGroupNonUniformFMul", 352)
val OpGroupNonUniformSMin = Code("OpGroupNonUniformSMin", 353)
val OpGroupNonUniformUMin = Code("OpGroupNonUniformUMin", 354)
val OpGroupNonUniformFMin = Code("OpGroupNonUniformFMin", 355)
val OpGroupNonUniformSMax = Code("OpGroupNonUniformSMax", 356)
val OpGroupNonUniformUMax = Code("OpGroupNonUniformUMax", 357)
val OpGroupNonUniformFMax = Code("OpGroupNonUniformFMax", 358)
val OpGroupNonUniformBitwiseAnd = Code("OpGroupNonUniformBitwiseAnd", 359)
val OpGroupNonUniformBitwiseOr = Code("OpGroupNonUniformBitwiseOr", 360)
val OpGroupNonUniformBitwiseXor = Code("OpGroupNonUniformBitwiseXor", 361)
val OpGroupNonUniformLogicalAnd = Code("OpGroupNonUniformLogicalAnd", 362)
val OpGroupNonUniformLogicalOr = Code("OpGroupNonUniformLogicalOr", 363)
val OpGroupNonUniformLogicalXor = Code("OpGroupNonUniformLogicalXor", 364)
val OpGroupNonUniformQuadBroadcast = Code("OpGroupNonUniformQuadBroadcast", 365)
val OpGroupNonUniformQuadSwap = Code("OpGroupNonUniformQuadSwap", 366)

object GlslOp:
val Round = Code("Round", 1)
val RoundEven = Code("RoundEven", 2)
Expand Down
9 changes: 8 additions & 1 deletion cyfra-compiler/src/main/scala/io/computenode/cyfra/spirv/SpirvConstants.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,5 +17,12 @@ private[cyfra] object SpirvConstants:
val GL_GLOBAL_INVOCATION_ID_REF = 5
val GL_WORKGROUP_SIZE_REF = 6
val DEBUG_PRINTF_REF = 7
val GL_LOCAL_INVOCATION_ID_REF = 8
val GL_LOCAL_INVOCATION_INDEX_REF = 9
val GL_WORKGROUP_ID_REF = 10
val GL_NUM_WORKGROUPS_REF = 11
val GL_SUBGROUP_ID_REF = 12
val GL_SUBGROUP_LOCAL_INVOCATION_ID_REF = 13
val GL_SUBGROUP_SIZE_REF = 14

val HEADER_REFS_TOP = 8
val HEADER_REFS_TOP = 15
42 changes: 41 additions & 1 deletion cyfra-compiler/src/main/scala/io/computenode/cyfra/spirv/SpirvTypes.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ private[cyfra] object SpirvTypes:

val Int32Tag = summon[Tag[Int32]]
val UInt32Tag = summon[Tag[UInt32]]
val Float16Tag = summon[Tag[Float16]]
val Float32Tag = summon[Tag[Float32]]
val GBooleanTag = summon[Tag[GBoolean]]
val Vec2TagWithoutArgs = summon[Tag[Vec2[?]]].tag.withoutArgs
Expand All @@ -22,6 +23,7 @@ private[cyfra] object SpirvTypes:

val LInt32Tag = Int32Tag.tag
val LUInt32Tag = UInt32Tag.tag
val LFloat16Tag = Float16Tag.tag
val LFloat32Tag = Float32Tag.tag
val LGBooleanTag = GBooleanTag.tag
val LVec2TagWithoutArgs = Vec2TagWithoutArgs
Expand All @@ -36,9 +38,38 @@ private[cyfra] object SpirvTypes:
type Vec3C[T <: Value] = Vec3[T]
type Vec4C[T <: Value] = Vec4[T]

/** Convert Float32 to Float16 (half precision) bits.
* Uses round-to-nearest-even rounding mode.
*/
def floatToFloat16(f: Float): Int = {
val bits = java.lang.Float.floatToIntBits(f)
val sign = (bits >>> 16) & 0x8000
val exponent = ((bits >>> 23) & 0xFF) - 127 + 15
val mantissa = bits & 0x007FFFFF

if (exponent <= 0) {
// Denormalized or zero
if (exponent < -10) {
sign // Zero
} else {
// Denormalized
val m = mantissa | 0x00800000
val shifted = m >>> (1 - exponent)
sign | (shifted >>> 13)
}
} else if (exponent >= 31) {
// Infinity or NaN
sign | 0x7C00 | (if (mantissa != 0) 0x200 else 0)
} else {
// Normalized
sign | (exponent << 10) | (mantissa >>> 13)
}
}

def scalarTypeDefInsn(tag: Tag[?], typeDefIndex: Int) = tag match
case Int32Tag => Instruction(Op.OpTypeInt, List(ResultRef(typeDefIndex), IntWord(32), IntWord(1)))
case UInt32Tag => Instruction(Op.OpTypeInt, List(ResultRef(typeDefIndex), IntWord(32), IntWord(0)))
case Float16Tag => Instruction(Op.OpTypeFloat, List(ResultRef(typeDefIndex), IntWord(16)))
case Float32Tag => Instruction(Op.OpTypeFloat, List(ResultRef(typeDefIndex), IntWord(32)))
case GBooleanTag => Instruction(Op.OpTypeBool, List(ResultRef(typeDefIndex)))

Expand All @@ -50,6 +81,7 @@ private[cyfra] object SpirvTypes:
def typeStride(tag: LightTypeTag): Int = tag match
case LInt32Tag => 4
case LUInt32Tag => 4
case LFloat16Tag => 2
case LFloat32Tag => 4
case LGBooleanTag => 4
case v if v <:< LVecTag =>
Expand All @@ -63,6 +95,14 @@ private[cyfra] object SpirvTypes:
IntWord(value.asInstanceOf[Int])
case t if t == UInt32Tag =>
IntWord(value.asInstanceOf[Int])
case t if t == Float16Tag =>
val fl = value match
case fl: Float => fl
case dl: Double => dl.toFloat
case il: Int => il.toFloat
// Convert Float32 to Float16 (half precision)
val f16Bits = floatToFloat16(fl)
Word(intToBytes(f16Bits & 0xFFFF).reverse.toArray)
case t if t == Float32Tag =>
val fl = value match
case fl: Float => fl
Expand All @@ -71,7 +111,7 @@ private[cyfra] object SpirvTypes:
Word(intToBytes(java.lang.Float.floatToIntBits(fl)).reverse.toArray)

def defineScalarTypes(types: List[Tag[?]], context: Context): (List[Words], Context) =
val basicTypes = List(Int32Tag, Float32Tag, UInt32Tag, GBooleanTag)
val basicTypes = List(Int32Tag, Float16Tag, Float32Tag, UInt32Tag, GBooleanTag)
(basicTypes ::: types).distinct.foldLeft((List[Words](), context)) { case ((words, ctx), valType) =>
val typeDefIndex = ctx.nextResultId
val code = List(
Expand Down
99 changes: 90 additions & 9 deletions cyfra-compiler/src/main/scala/io/computenode/cyfra/spirv/compilers/DSLCompiler.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ import io.computenode.cyfra.*
import io.computenode.cyfra.dsl.*
import io.computenode.cyfra.dsl.Expression.E
import io.computenode.cyfra.dsl.Value.Scalar
import io.computenode.cyfra.dsl.binding.{GBinding, GBuffer, GUniform, WriteBuffer, WriteUniform}
import io.computenode.cyfra.dsl.binding.{GBinding, GBuffer, GShared, GUniform, ReadShared, WriteBuffer, WriteShared, WriteUniform}
import io.computenode.cyfra.dsl.gio.GIO
import io.computenode.cyfra.dsl.struct.GStruct.*
import io.computenode.cyfra.dsl.struct.GStructSchema
Expand Down Expand Up @@ -34,7 +34,7 @@ private[cyfra] object DSLCompiler:
getAllExprsFlattened(tail, getAllExprsFlattened(v.tree, visitDetached) ::: acc, visitDetached)
case GIO.FlatMap(v, n) :: tail =>
getAllExprsFlattened(v :: n :: tail, acc, visitDetached)
case GIO.Repeat(n, gio) :: tail =>
case GIO.Repeat(n, gio, _) :: tail =>
val nAllExprs = getAllExprsFlattened(n.tree, visitDetached)
getAllExprsFlattened(gio :: tail, nAllExprs ::: acc, visitDetached)
case WriteBuffer(_, index, value) :: tail =>
Expand All @@ -47,6 +47,16 @@ private[cyfra] object DSLCompiler:
case GIO.Printf(_, args*) :: tail =>
val argsAllExprs = args.flatMap(a => getAllExprsFlattened(a.tree, visitDetached)).toList
getAllExprsFlattened(tail, argsAllExprs ::: acc, visitDetached)
case GIO.WorkgroupBarrier :: tail =>
getAllExprsFlattened(tail, acc, visitDetached)
case WriteShared(_, index, value) :: tail =>
val indexAllExprs = getAllExprsFlattened(index.tree, visitDetached)
val valueAllExprs = getAllExprsFlattened(value.tree, visitDetached)
getAllExprsFlattened(tail, indexAllExprs ::: valueAllExprs ::: acc, visitDetached)
case GIO.FoldRepeat(n, init, body, _, _) :: tail =>
val nAllExprs = getAllExprsFlattened(n.tree, visitDetached)
val initAllExprs = getAllExprsFlattened(init.tree, visitDetached)
getAllExprsFlattened(body :: tail, nAllExprs ::: initAllExprs ::: acc, visitDetached)

// TODO: Not traverse same fn scopes for each fn call
private def getAllExprsFlattened(root: E[?], visitDetached: Boolean): List[E[?]] =
Expand All @@ -71,22 +81,82 @@ private[cyfra] object DSLCompiler:
allScopesCache(root.treeid) = result
result

private def getAllShared(pending: List[GIO[?]], acc: Map[Int, GShared[?]]): Map[Int, GShared[?]] =
pending match
case Nil => acc
case GIO.FlatMap(v, n) :: tail =>
getAllShared(v :: n :: tail, acc)
case GIO.Repeat(_, gio, _) :: tail =>
getAllShared(gio :: tail, acc)
case GIO.FoldRepeat(_, _, gio, _, _) :: tail =>
getAllShared(gio :: tail, acc)
case WriteShared(buffer, _, _) :: tail =>
val impl = buffer.asInstanceOf[GShared.GSharedImpl[?]]
getAllShared(tail, acc + (impl.sharedId -> buffer))
case _ :: tail => getAllShared(tail, acc)

private def getAllSharedFromExprs(exprs: List[E[?]], acc: Map[Int, GShared[?]]): Map[Int, GShared[?]] =
exprs.foldLeft(acc):
case (a, ReadShared(buffer, _)) =>
val impl = buffer.asInstanceOf[GShared.GSharedImpl[?]]
a + (impl.sharedId -> buffer)
case (a, _) => a

private def createSharedVariables(sharedBuffers: Map[Int, GShared[?]], ctx: Context): (List[Words], Context) =
sharedBuffers.foldLeft((List.empty[Words], ctx)):
case ((insnsAcc, currentCtx), (sharedId, buffer)) =>
val elementTypeRef = currentCtx.valueTypeMap(buffer.tag.tag)
val arraySizeConstRef = currentCtx.constRefs.getOrElse(
(Int32Tag, buffer.size),
throw new IllegalStateException(s"Missing constant for shared array size ${buffer.size}"),
)

// SPIR-V shared memory structure:
// 1. Array type: OpTypeArray %arrayType %elementType %size
// 2. Pointer to array: OpTypePointer %ptrArrayType Workgroup %arrayType
// 3. Variable: OpVariable %ptrArrayType %var Workgroup
// 4. Pointer to element: OpTypePointer %ptrElemType Workgroup %elementType (for OpAccessChain)
val arrayTypeRef = currentCtx.nextResultId
val ptrArrayTypeRef = currentCtx.nextResultId + 1
val varRef = currentCtx.nextResultId + 2
val ptrElemTypeRef = currentCtx.nextResultId + 3

val insns = List(
Instruction(Op.OpTypeArray, List(ResultRef(arrayTypeRef), ResultRef(elementTypeRef), ResultRef(arraySizeConstRef))),
Instruction(Op.OpTypePointer, List(ResultRef(ptrArrayTypeRef), StorageClass.Workgroup, ResultRef(arrayTypeRef))),
Instruction(Op.OpVariable, List(ResultRef(ptrArrayTypeRef), ResultRef(varRef), StorageClass.Workgroup)),
Instruction(Op.OpTypePointer, List(ResultRef(ptrElemTypeRef), StorageClass.Workgroup, ResultRef(elementTypeRef))),
)

val block = SharedBlock(arrayTypeRef, varRef, ptrElemTypeRef)
val newCtx = currentCtx.copy(
nextResultId = currentCtx.nextResultId + 4,
sharedVarRefs = currentCtx.sharedVarRefs + (sharedId -> block),
workgroupPointerMap = currentCtx.workgroupPointerMap + (elementTypeRef -> ptrElemTypeRef),
)
(insnsAcc ::: insns, newCtx)

// So far only used for printf
private def getAllStrings(pending: List[GIO[?]], acc: Set[String]): Set[String] =
pending match
case Nil => acc
case GIO.FlatMap(v, n) :: tail =>
getAllStrings(v :: n :: tail, acc)
case GIO.Repeat(_, gio) :: tail =>
case GIO.Repeat(_, gio, _) :: tail =>
getAllStrings(gio :: tail, acc)
case GIO.Printf(format, _*) :: tail =>
getAllStrings(tail, acc + format)
case _ :: tail => getAllStrings(tail, acc)

def compile(bodyIo: GIO[?], bindings: List[GBinding[?]]): ByteBuffer =
def compile(bodyIo: GIO[?], bindings: List[GBinding[?]], workgroupSize: (Int, Int, Int) = (256, 1, 1)): ByteBuffer =
val allExprs = getAllExprsFlattened(List(bodyIo), Nil, visitDetached = true)
val typesInCode = allExprs.map(_.tag).distinct
val allTypes = (typesInCode ::: bindings.map(_.tag)).distinct

val sharedFromGio = getAllShared(List(bodyIo), Map.empty)
val sharedFromExprs = getAllSharedFromExprs(allExprs, sharedFromGio)
val sharedTypes = sharedFromExprs.values.map(_.tag).toList

val allTypes = (typesInCode ::: bindings.map(_.tag) ::: sharedTypes).distinct
def scalarTypes = allTypes.filter(_.tag <:< summon[Tag[Scalar]].tag)
val (typeDefs, typedContext) = defineScalarTypes(scalarTypes, Context.initialContext)
val allStrings = getAllStrings(List(bodyIo), Set.empty)
Expand All @@ -108,17 +178,28 @@ private[cyfra] object DSLCompiler:
val (decorations, uniformDefs, uniformContext) = initAndDecorateBuffers(buffersWithIndices, structNamesCtx)
val (uniformStructDecorations, uniformStructInsns, uniformStructContext) = createAndInitUniformBlocks(uniformsWithIndices, uniformContext)
val blockNames = getBlockNames(uniformContext, uniforms)
val (inputDefs, inputContext) = createInvocationId(uniformStructContext)
val (inputDefs, inputContext) = createInvocationId(uniformStructContext, workgroupSize)

val sharedSizeConsts = sharedFromExprs.values.map(s => (Int32Tag, s.size)).toList
val (constDefs, constCtx) = defineConstants(allExprs, inputContext)
val (varDefs, varCtx) = defineVarNames(constCtx)
val (sharedConstDefs, constCtxWithShared) = sharedSizeConsts.foldLeft((List.empty[Words], constCtx)):
case ((insnsAcc, ctx), const) if ctx.constRefs.contains(const) => (insnsAcc, ctx)
case ((insnsAcc, ctx), const) =>
val insn = Instruction(Op.OpConstant, List(ResultRef(ctx.valueTypeMap(const._1.tag)), ResultRef(ctx.nextResultId), IntWord(const._2)))
val newCtx = ctx.copy(constRefs = ctx.constRefs + (const -> ctx.nextResultId), nextResultId = ctx.nextResultId + 1)
(insnsAcc :+ insn, newCtx)

val (sharedDefs, ctxWithShared) = createSharedVariables(sharedFromExprs, constCtxWithShared)

val (varDefs, varCtx) = defineVarNames(ctxWithShared)
val (main, ctxAfterMain) = compileMain(bodyIo, varCtx)
val (fnTypeDefs, fnDefs, ctxWithFnDefs) = compileFunctions(ctxAfterMain)
val nameDecorations = getNameDecorations(ctxWithFnDefs)

val code: List[Words] =
SpirvProgramCompiler.headers ::: stringDefs ::: blockNames ::: nameDecorations ::: structNames ::: SpirvProgramCompiler.workgroupDecorations :::
SpirvProgramCompiler.headers(workgroupSize) ::: stringDefs ::: blockNames ::: nameDecorations ::: structNames ::: SpirvProgramCompiler.workgroupDecorations :::
decorations ::: uniformStructDecorations ::: typeDefs ::: structDefs ::: fnTypeDefs ::: uniformDefs ::: uniformStructInsns ::: inputDefs :::
constDefs ::: varDefs ::: main ::: fnDefs
constDefs ::: sharedConstDefs ::: sharedDefs ::: varDefs ::: main ::: fnDefs

val fullCode = code.map:
case WordVariable(name) if name == BOUND_VARIABLE => IntWord(ctxWithFnDefs.nextResultId)
Expand Down
Loading