Implementation:Ollama Ollama Imagegen MLX Compile

Knowledge Sources	Ollama
Domains	Image Generation, Native Bindings
Last Updated	2025-02-15 00:00 GMT

Overview

Provides Go function compilation into optimized MLX closures, enabling graph-level optimization for repeated operations.

Description

The compile.go file implements MLX function compilation via CGo. The Compile function takes a Go ClosureFunc (func([]*Array) []*Array), wraps it as an mlx_closure via a CGo callback, and compiles it into an optimized MLX execution graph. The compiled closure runs the optimized graph on subsequent calls without creating Go intermediate arrays. A CompiledFunc struct holds both the original and compiled closures, with Call/CallEval methods for invocation. The file manages closure callbacks carefully: inClosureCallback flag disables GC cleanup during callbacks (since MLX owns the arrays), borrowArray wraps C arrays without finalizers, and goClosureDestructor cleans up Go handles when MLX frees the closure.

Usage

Used for performance-critical operations like SwiGLU activation in GPT-OSS, where graph compilation eliminates per-call overhead.

Code Reference

Source Location

• Repository: Ollama
• File: x/imagegen/mlx/compile.go
• Lines: 1-173

Signature

type ClosureFunc func(inputs []*Array) []*Array

type CompiledFunc struct {
	closure  C.mlx_closure
	compiled C.mlx_closure
}

func Compile(fn ClosureFunc) *CompiledFunc
func CompileShapeless(fn ClosureFunc, shapeless bool) *CompiledFunc
func (cf *CompiledFunc) Call(inputs ...*Array) []*Array
func (cf *CompiledFunc) CallEval(inputs ...*Array) []*Array
func (cf *CompiledFunc) Free()
func InClosureCallback() bool

Import

import "github.com/ollama/ollama/x/imagegen/mlx"

I/O Contract

Inputs

Name	Type	Required	Description
fn	ClosureFunc	Yes	Go function to compile into MLX graph
shapeless	bool	No	If true, compiled graph works for any input shape

Outputs

Name	Type	Description
*CompiledFunc	*CompiledFunc	Compiled function with Call/Free methods

Usage Examples

compiled := mlx.Compile(func(inputs []*mlx.Array) []*mlx.Array {
    a, b := inputs[0], inputs[1]
    c := mlx.Add(a, b)
    d := mlx.Mul(c, c)
    return []*mlx.Array{d}
})
defer compiled.Free()

result := compiled.Call(x, y)[0]
mlx.Eval(result)

Related Pages

• Principle:Ollama_Ollama_ImageGeneration