include/brotensor/jit/trace.h - GCC Code Coverage Report

Directory:	brotensor/
Coverage:	low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%

	Coverage	Exec / Excl / Total
Lines:	100.0%	4 / 0 / 4
Functions:	100.0%	4 / 0 / 4
Branches:	-%	0 / 0 / 0

    include/brotensor/jit/trace.h
    
        Line
        Branch
        Exec
        Source
      
        #pragma once
      
        // ─── brotensor::jit — Automatic Tracing JIT Interface ────────────────────────
      
        //
      
        // Dynamic expression and operation sequence tracing powered by Brass.
      
        // Automatically captures mathematical expressions and operation patterns,
      
        // analyzes tensor liveness, and emits fused AVX2/FMA machine code on CPU,
      
        // cooperative PTX kernels on NVIDIA GPUs, or SPIR-V compute kernels on Vulkan
      
        // (brass's SPIR-V target; a Vulkan trace with no fusion is replayed op by op).
      
        #include <brotensor/tensor.h>
      
        #include <cstdint>
      
        #include <cstddef>
      
        #include <memory>
      
        #include <stdexcept>
      
        namespace brotensor::jit {
      
        class TraceHandleImpl;
      
        class TraceHandle {
      
        public:
      
            TraceHandle() = default;
      
        22
            explicit TraceHandle(std::shared_ptr<TraceHandleImpl> impl) : impl_(std::move(impl)) {}
      
            void execute();
      
            uint64_t signature_hash() const;
      
            bool is_cuda() const;
      
            size_t node_count() const;
      
            double last_duration_us() const;
      
            bool is_cache_hit() const;
      
            // How many kernel launches one execute() issues. A fully fused trace is 1;
      
            // a trace the compiler could only partly fuse reports what it actually
      
            // costs, which is the number the benchmarks compare against the eager op
      
            // sequence's launch count.
      
            size_t launch_count() const;
      
            // Wall-clock cost of producing this trace's code, in microseconds: PTX
      
            // emission plus the driver's PTX->SASS compile (CUDA), or MIR, SPIR-V and
      
            // pipeline creation (Vulkan). Zero on a cache hit.
      
            double compile_us() const;
      
            // Human-readable name of the fusion the compiler chose, for benchmark and
      
            // test output ("elementwise-vec8-bf16", "row-norm", "residual-rmsnorm", …).
      
            const char* fusion_name() const;
      
            explicit operator bool() const noexcept { return impl_ != nullptr; }
      
        private:
      
            std::shared_ptr<TraceHandleImpl> impl_;
      
        };
      
        // ── What a traced op returns ──────────────────────────────────────────────
      
        //
      
        // Between begin_trace() and end_trace(), the operators and functions below do
      
        // not compute anything and do not allocate anything. Each returns a *symbolic*
      
        // Tensor: a name for one node of the trace, carrying the right shape, dtype
      
        // and device but no buffer. The fused kernel evaluates the whole expression in
      
        // registers, so a full-size buffer per intermediate would be pure waste — at
      
        // VAE decode width it is hundreds of megabytes per trace for contents nothing
      
        // ever reads.
      
        //
      
        // Three consequences for the caller:
      
        //
      
        //   * A symbolic Tensor cannot be copied. It is a value in an expression, not
      
        //     a tensor; the copy ctor throws. Use it, move it, or store() it.
      
        //
      
        //   * end_trace() gives a real buffer to every symbolic Tensor the caller is
      
        //     still holding whose value the trace actually produces — in
      
        //     `Tensor out = (a*b+c)*silu(d); auto h = end_trace();`, `out` owns a real
      
        //     buffer afterwards and `h.execute()` refills it, exactly as before.
      
        //
      
        //   * A symbolic Tensor whose value the fused kernel only consumes internally
      
        //     — `Tensor ln = layernorm(x); Tensor out = ln * g;`, where nothing else
      
        //     reads `ln` — comes back from end_trace() as an *empty* tensor. Its
      
        //     contents were never written to memory. (Before, it came back holding an
      
        //     allocated buffer of uninitialised garbage, so no correct code can be
      
        //     relying on it.)
      
        // Starts tracing thread-local operations.
      
        void begin_trace();
      
        // Stops tracing, compiles the captured DAG (or retrieves from cache), and
      
        // returns a TraceHandle. Materialises the live-at-end intermediates the caller
      
        // still holds, as described above.
      
        //
      
        // If compilation throws — the expression has no fusion on this backend — the
      
        // trace has already been closed and those tensors have already been given
      
        // buffers; they hold undefined contents and the caller is expected to run its
      
        // eager path over them. abort_trace() afterwards is harmless.
      
        TraceHandle end_trace();
      
        // Discards a trace in progress without compiling it. For a caller that has to
      
        // unwind out of a traced region — an unsupported expression throws from
      
        // end_trace(), and the thread's trace context has to be left clean for the
      
        // next attempt.
      
        //
      
        // Every symbolic Tensor the caller still holds is neutralised to an empty
      
        // tensor (no buffer, zero rows and cols, jit_slot cleared) rather than
      
        // materialised: the abort path exists so the caller can recompute the whole
      
        // expression eagerly, and allocating buffers it is about to overwrite would
      
        // cost exactly what tracing is here to avoid. A caller that wants values back
      
        // assigns over those tensors on its eager path, which is what
      
        // brodiffusion::detail::try_fused does.
      
        void abort_trace();
      
        // Returns true if tracing is currently active on this thread.
      
        bool is_tracing();
      
        // RAII helper for tracing scopes.
      
        class TraceScope {
      
        public:
      
            TraceScope() { begin_trace(); }
      
            ~TraceScope() {
      
                if (is_tracing()) {
      
                    try { end_trace(); } catch (...) {}
      
                }
      
            }
      
            TraceHandle end() { return end_trace(); }
      
            TraceScope(const TraceScope&) = delete;
      
            TraceScope& operator=(const TraceScope&) = delete;
      
        };
      
        // ── Overloaded Operators for Tracing and Eager Execution ──────────────────────
      
        Tensor operator+(const Tensor& a, const Tensor& b);
      
        Tensor operator+(const Tensor& a, float s);
      
        1
        inline Tensor operator+(float s, const Tensor& a) { return a + s; }
      
        inline Tensor operator+(const Tensor& a, int s) { return a + static_cast<float>(s); }
      
        inline Tensor operator+(int s, const Tensor& a) { return a + static_cast<float>(s); }
      
        Tensor operator-(const Tensor& a, const Tensor& b);
      
        Tensor operator-(const Tensor& a, float s);
      
        Tensor operator-(float s, const Tensor& a);
      
        inline Tensor operator-(const Tensor& a, int s) { return a - static_cast<float>(s); }
      
        inline Tensor operator-(int s, const Tensor& a) { return static_cast<float>(s) - a; }
      
        Tensor operator*(const Tensor& a, const Tensor& b);
      
        Tensor operator*(const Tensor& a, float s);
      
        inline Tensor operator*(float s, const Tensor& a) { return a * s; }
      
        inline Tensor operator*(const Tensor& a, int s) { return a * static_cast<float>(s); }
      
        inline Tensor operator*(int s, const Tensor& a) { return a * static_cast<float>(s); }
      
        Tensor operator/(const Tensor& a, const Tensor& b);
      
        Tensor operator/(const Tensor& a, float s);
      
        inline Tensor operator/(const Tensor& a, int s) { return a / static_cast<float>(s); }
      
        Tensor& operator+=(Tensor& a, const Tensor& b);
      
        Tensor& operator+=(Tensor& a, float s);
      
        inline Tensor& operator+=(Tensor& a, int s) { return a += static_cast<float>(s); }
      
        Tensor& operator*=(Tensor& a, const Tensor& b);
      
        Tensor& operator*=(Tensor& a, float s);
      
        inline Tensor& operator*=(Tensor& a, int s) { return a *= static_cast<float>(s); }
      
        // ── Functional Activations & Normalizations ───────────────────────────────────
      
        Tensor silu(const Tensor& a);
      
        Tensor gelu(const Tensor& a);
      
        Tensor relu(const Tensor& a);
      
        Tensor tanh(const Tensor& a);
      
        Tensor sigmoid(const Tensor& a);
      
        Tensor rms_norm(const Tensor& x, const Tensor& gamma = Tensor(), float eps = 1e-5f);
      
        Tensor layernorm(const Tensor& x, const Tensor& gamma = Tensor(), const Tensor& beta = Tensor(), float eps = 1e-5f);
      
        Tensor modulate(const Tensor& x, const Tensor& scale, const Tensor& shift);
      
        // Elementwise conditional selection: cond != 0 ? a : b (supports full, row,
      
        // col, and scalar broadcast operands in trace mode).
      
        Tensor where(const Tensor& cond, const Tensor& a, const Tensor& b);
      
        1
        inline Tensor select(const Tensor& cond, const Tensor& a, const Tensor& b) {
      
        1
            return where(cond, a, b);
      
        }
      
        // Functional elementwise clamp to [lo, hi]. Note: in-place `clamp(Tensor& y, lo, hi)`
      
        // from <brotensor/ops/elementwise.h> also records a Clamp node when tracing is active.
      
        Tensor clamp(const Tensor& x, float lo, float hi);
      
        inline Tensor clip(const Tensor& x, float lo, float hi) { return clamp(x, lo, hi); }
      
        // Functional dtype conversion to `out_dtype` (FP32 / FP16 / BF16). Note: 3-arg
      
        // `cast(src, dst, out_dtype)` also records a Cast + store when tracing is active.
      
        Tensor cast(const Tensor& x, Dtype out_dtype);
      
        // Writes `src` into `dst`'s existing buffer. Under a trace this is what makes
      
        // an expression land somewhere the caller already owns — a preallocated
      
        // scratch tensor, or a row view into a larger one — instead of in a buffer the
      
        // tracer allocated and the caller would have to copy out of. Outside a trace
      
        // it is an ordinary device-to-device copy.
      
        void store(Tensor& dst, const Tensor& src);
      
        } // namespace brotensor::jit
      
        namespace brotensor {
      
        using jit::begin_trace;
      
        using jit::end_trace;
      
        using jit::abort_trace;
      
        using jit::is_tracing;
      
        using jit::store;
      
        using jit::TraceScope;
      
        using jit::TraceHandle;
      
        using jit::operator+;
      
        using jit::operator-;
      
        using jit::operator*;
      
        using jit::operator/;
      
        using jit::operator+=;
      
        using jit::operator*=;
      
        using jit::silu;
      
        using jit::gelu;
      
        using jit::relu;
      
        using jit::tanh;
      
        using jit::sigmoid;
      
        using jit::rms_norm;
      
        using jit::layernorm;
      
        using jit::modulate;
      
        using jit::where;
      
        using jit::select;
      
        using jit::clamp;
      
        using jit::clip;
      
        using jit::cast;
      
        } // namespace brotensor

Line	Exec	Source
1		#pragma once
2
3		// ─── brotensor::jit — Automatic Tracing JIT Interface ────────────────────────
4		//
5		// Dynamic expression and operation sequence tracing powered by Brass.
6		// Automatically captures mathematical expressions and operation patterns,
7		// analyzes tensor liveness, and emits fused AVX2/FMA machine code on CPU,
8		// cooperative PTX kernels on NVIDIA GPUs, or SPIR-V compute kernels on Vulkan
9		// (brass's SPIR-V target; a Vulkan trace with no fusion is replayed op by op).
10
11		#include <brotensor/tensor.h>
12		#include <cstdint>
13		#include <cstddef>
14		#include <memory>
15		#include <stdexcept>
16
17		namespace brotensor::jit {
18
19		class TraceHandleImpl;
20
21		class TraceHandle {
22		public:
23		TraceHandle() = default;
24	22	explicit TraceHandle(std::shared_ptr<TraceHandleImpl> impl) : impl_(std::move(impl)) {}
25
26		void execute();
27		uint64_t signature_hash() const;
28		bool is_cuda() const;
29		size_t node_count() const;
30		double last_duration_us() const;
31		bool is_cache_hit() const;
32
33		// How many kernel launches one execute() issues. A fully fused trace is 1;
34		// a trace the compiler could only partly fuse reports what it actually
35		// costs, which is the number the benchmarks compare against the eager op
36		// sequence's launch count.
37		size_t launch_count() const;
38
39		// Wall-clock cost of producing this trace's code, in microseconds: PTX
40		// emission plus the driver's PTX->SASS compile (CUDA), or MIR, SPIR-V and
41		// pipeline creation (Vulkan). Zero on a cache hit.
42		double compile_us() const;
43
44		// Human-readable name of the fusion the compiler chose, for benchmark and
45		// test output ("elementwise-vec8-bf16", "row-norm", "residual-rmsnorm", …).
46		const char* fusion_name() const;
47
48		explicit operator bool() const noexcept { return impl_ != nullptr; }
49
50		private:
51		std::shared_ptr<TraceHandleImpl> impl_;
52		};
53
54		// ── What a traced op returns ──────────────────────────────────────────────
55		//
56		// Between begin_trace() and end_trace(), the operators and functions below do
57		// not compute anything and do not allocate anything. Each returns a symbolic
58		// Tensor: a name for one node of the trace, carrying the right shape, dtype
59		// and device but no buffer. The fused kernel evaluates the whole expression in
60		// registers, so a full-size buffer per intermediate would be pure waste — at
61		// VAE decode width it is hundreds of megabytes per trace for contents nothing
62		// ever reads.
63		//
64		// Three consequences for the caller:
65		//
66		// * A symbolic Tensor cannot be copied. It is a value in an expression, not
67		// a tensor; the copy ctor throws. Use it, move it, or store() it.
68		//
69		// * end_trace() gives a real buffer to every symbolic Tensor the caller is
70		// still holding whose value the trace actually produces — in
71		// `Tensor out = (ab+c)silu(d); auto h = end_trace();`, `out` owns a real
72		// buffer afterwards and `h.execute()` refills it, exactly as before.
73		//
74		// * A symbolic Tensor whose value the fused kernel only consumes internally
75		// — `Tensor ln = layernorm(x); Tensor out = ln * g;`, where nothing else
76		// reads `ln` — comes back from end_trace() as an empty tensor. Its
77		// contents were never written to memory. (Before, it came back holding an
78		// allocated buffer of uninitialised garbage, so no correct code can be
79		// relying on it.)
80
81		// Starts tracing thread-local operations.
82		void begin_trace();
83
84		// Stops tracing, compiles the captured DAG (or retrieves from cache), and
85		// returns a TraceHandle. Materialises the live-at-end intermediates the caller
86		// still holds, as described above.
87		//
88		// If compilation throws — the expression has no fusion on this backend — the
89		// trace has already been closed and those tensors have already been given
90		// buffers; they hold undefined contents and the caller is expected to run its
91		// eager path over them. abort_trace() afterwards is harmless.
92		TraceHandle end_trace();
93
94		// Discards a trace in progress without compiling it. For a caller that has to
95		// unwind out of a traced region — an unsupported expression throws from
96		// end_trace(), and the thread's trace context has to be left clean for the
97		// next attempt.
98		//
99		// Every symbolic Tensor the caller still holds is neutralised to an empty
100		// tensor (no buffer, zero rows and cols, jit_slot cleared) rather than
101		// materialised: the abort path exists so the caller can recompute the whole
102		// expression eagerly, and allocating buffers it is about to overwrite would
103		// cost exactly what tracing is here to avoid. A caller that wants values back
104		// assigns over those tensors on its eager path, which is what
105		// brodiffusion::detail::try_fused does.
106		void abort_trace();
107
108		// Returns true if tracing is currently active on this thread.
109		bool is_tracing();
110
111		// RAII helper for tracing scopes.
112		class TraceScope {
113		public:
114		TraceScope() { begin_trace(); }
115		~TraceScope() {
116		if (is_tracing()) {
117		try { end_trace(); } catch (...) {}
118		}
119		}
120		TraceHandle end() { return end_trace(); }
121
122		TraceScope(const TraceScope&) = delete;
123		TraceScope& operator=(const TraceScope&) = delete;
124		};
125
126		// ── Overloaded Operators for Tracing and Eager Execution ──────────────────────
127
128		Tensor operator+(const Tensor& a, const Tensor& b);
129		Tensor operator+(const Tensor& a, float s);
130	1	inline Tensor operator+(float s, const Tensor& a) { return a + s; }
131		inline Tensor operator+(const Tensor& a, int s) { return a + static_cast<float>(s); }
132		inline Tensor operator+(int s, const Tensor& a) { return a + static_cast<float>(s); }
133
134		Tensor operator-(const Tensor& a, const Tensor& b);
135		Tensor operator-(const Tensor& a, float s);
136		Tensor operator-(float s, const Tensor& a);
137		inline Tensor operator-(const Tensor& a, int s) { return a - static_cast<float>(s); }
138		inline Tensor operator-(int s, const Tensor& a) { return static_cast<float>(s) - a; }
139
140		Tensor operator*(const Tensor& a, const Tensor& b);
141		Tensor operator*(const Tensor& a, float s);
142		inline Tensor operator(float s, const Tensor& a) { return a s; }
143		inline Tensor operator(const Tensor& a, int s) { return a static_cast<float>(s); }
144		inline Tensor operator(int s, const Tensor& a) { return a static_cast<float>(s); }
145
146		Tensor operator/(const Tensor& a, const Tensor& b);
147		Tensor operator/(const Tensor& a, float s);
148		inline Tensor operator/(const Tensor& a, int s) { return a / static_cast<float>(s); }
149
150		Tensor& operator+=(Tensor& a, const Tensor& b);
151		Tensor& operator+=(Tensor& a, float s);
152		inline Tensor& operator+=(Tensor& a, int s) { return a += static_cast<float>(s); }
153
154		Tensor& operator*=(Tensor& a, const Tensor& b);
155		Tensor& operator*=(Tensor& a, float s);
156		inline Tensor& operator=(Tensor& a, int s) { return a = static_cast<float>(s); }
157
158		// ── Functional Activations & Normalizations ───────────────────────────────────
159
160		Tensor silu(const Tensor& a);
161		Tensor gelu(const Tensor& a);
162		Tensor relu(const Tensor& a);
163		Tensor tanh(const Tensor& a);
164		Tensor sigmoid(const Tensor& a);
165
166		Tensor rms_norm(const Tensor& x, const Tensor& gamma = Tensor(), float eps = 1e-5f);
167		Tensor layernorm(const Tensor& x, const Tensor& gamma = Tensor(), const Tensor& beta = Tensor(), float eps = 1e-5f);
168		Tensor modulate(const Tensor& x, const Tensor& scale, const Tensor& shift);
169
170		// Elementwise conditional selection: cond != 0 ? a : b (supports full, row,
171		// col, and scalar broadcast operands in trace mode).
172		Tensor where(const Tensor& cond, const Tensor& a, const Tensor& b);
173	1	inline Tensor select(const Tensor& cond, const Tensor& a, const Tensor& b) {
174	1	return where(cond, a, b);
175		}
176
177		// Functional elementwise clamp to [lo, hi]. Note: in-place `clamp(Tensor& y, lo, hi)`
178		// from <brotensor/ops/elementwise.h> also records a Clamp node when tracing is active.
179		Tensor clamp(const Tensor& x, float lo, float hi);
180		inline Tensor clip(const Tensor& x, float lo, float hi) { return clamp(x, lo, hi); }
181
182		// Functional dtype conversion to `out_dtype` (FP32 / FP16 / BF16). Note: 3-arg
183		// `cast(src, dst, out_dtype)` also records a Cast + store when tracing is active.
184		Tensor cast(const Tensor& x, Dtype out_dtype);
185
186		// Writes `src` into `dst`'s existing buffer. Under a trace this is what makes
187		// an expression land somewhere the caller already owns — a preallocated
188		// scratch tensor, or a row view into a larger one — instead of in a buffer the
189		// tracer allocated and the caller would have to copy out of. Outside a trace
190		// it is an ordinary device-to-device copy.
191		void store(Tensor& dst, const Tensor& src);
192
193		} // namespace brotensor::jit
194
195		namespace brotensor {
196		using jit::begin_trace;
197		using jit::end_trace;
198		using jit::abort_trace;
199		using jit::is_tracing;
200		using jit::store;
201		using jit::TraceScope;
202		using jit::TraceHandle;
203
204		using jit::operator+;
205		using jit::operator-;
206		using jit::operator*;
207		using jit::operator/;
208		using jit::operator+=;
209		using jit::operator*=;
210
211		using jit::silu;
212		using jit::gelu;
213		using jit::relu;
214		using jit::tanh;
215		using jit::sigmoid;
216		using jit::rms_norm;
217		using jit::layernorm;
218		using jit::modulate;
219		using jit::where;
220		using jit::select;
221		using jit::clamp;
222		using jit::clip;
223		using jit::cast;
224		} // namespace brotensor
225