I am implementing a sequential Model struct in Zig that manages layers and tensors with an autograd feature. However, when I test my model, I encounter a panic: reached unreachable code error.
Error Output (Truncated)
Freeing Tensor x.Wt+b
Freeing Tensor
Freeing Tensor ll.w
Freeing Tensor ll.b
Freeing Tensor thread 125464 panic: reached unreachable code
/home/am/zig/lib/std/posix.zig:1263:23: 0x107b73e in write (test)
.FAULT => unreachable,
^
...
/home/am/Zignet/src/tensor.zig:61:24: 0x1045d5c in deinit (test)
std.debug.print("Freeing Tensor {s}\n", .{self.label});
^
/home/am/Zignet/src/layers.zig:107:28: 0x104d31b in deinit (test)
self.weights.deinit();
^
/home/am/Zignet/src/sequential.zig:52:25: 0x1048d3a in deinit (test)
layer.deinit();
^
Code Snippet
Here is a simplified version of my Model struct:
const std = @import("std");
const Tensor = @import("tensor.zig").Tensor;
const Layer = @import("layers.zig").Layer;
const Linear = @import("layers.zig").Linear;
pub const Model = struct {
allocator: std.mem.Allocator,
layers: std.ArrayList(*Layer),
pub fn init(allocator: std.mem.Allocator) !Model {
const model = Model{
.allocator = allocator,
.layers = std.ArrayList(*Layer).init(allocator),
};
return model;
}
pub fn add_layer(self: *Model, layer: *Layer) !void {
try self.layers.append(layer);
}
pub fn forward(self: *Model, input: *Tensor) !*Tensor {
var output = input;
for (self.layers.items) |layer| {
output = try layer.forward(output);
}
return output;
}
pub fn zero_grad(self: *Model) void {
for (self.layers.items) |layer| {
layer.zero_grad();
}
}
pub fn parameters(self: *Model) !std.ArrayList(*Tensor) {
var params = std.ArrayList(*Tensor).init(self.allocator);
for (self.layers.items) |layer| {
switch (layer.*) {
.Linear => |linear| {
try params.append(linear.weights);
try params.append(linear.biases);
},
}
}
return params;
}
pub fn deinit(self: *Model) void {
for (self.layers.items) |layer| {
layer.deinit();
}
self.layers.deinit();
}
};
When I test layers.zig and tensor.zig separately, all tests pass. However, when using Model, the test crashes. My optimizer, which retrieves parameters from Model, is also affected.
Suspected Issue
Based on debugging, I suspect the issue is related to the deinit() order. It seems that layers are freed, but Model later tries to free them again, causing a panic.
Questions:
- Is my
deinit() implementation causing a double free?
- Should I restructure the
Model to better manage memory in Zig?
- How can I ensure that layers are properly freed without causing errors?
My full code is available here: GitHub Repository.
I am using Zig 0.14.0-dev.2647+5322459a0 on Ubuntu 22.04.5 LTS. Any help is appreciated since I am new to Zig and unsure if this Python-friendly structure is appropriate for Zig.
It looks like it might be a double free. In your test you defer deinit the layer and then add it to your model. And you are also deiniting your model. Line 65 and 74 of sequential.zig.
I agree with this observation. But when I remove the deinit of the layer, I get a memory leak. Here is my current test
const expect = std.testing.expect;
test "Model and Layer functionality" {
const allocator = std.testing.allocator;
// Create a Linear layer
const input_size: usize = 2;
const output_size: usize = 2;
var linear_layer = try Linear.init(allocator, input_size, output_size, null);
// Wrap the Linear layer in a Layer union
var linear_layer_wrapper = Layer{ .Linear = &linear_layer };
// Initialize a model
var model = try Model.init(allocator);
defer model.deinit();
try model.add_layer(&linear_layer_wrapper);
// Create a dummy input tensor
const input_shape = &[_]usize{ 1, input_size };
const input = try Tensor.from_value(allocator, input_shape, dtypes.DataType.f32, 1.0, true);
defer input.deinit();
// Perform a forward pass
const output = try model.forward(input);
defer output.deinit();
// Check the output shape
try expect(output.shape[0] == 1);
try expect(output.shape[1] == output_size);
// Check the parameters
const params = try model.parameters();
defer params.deinit();
// Linear layer has weights and biases
try expect(params.items.len == 2);
}
And this causes this
[gpa] (err): memory address 0x73dd4b027700 leaked:
/home/am/work/zignet/src/tensor.zig:154:44: 0x1054666 in from_ndarray (test)
const tensor = try allocator.create(Tensor);
^
/home/am/work/zignet/src/autograd.zig:424:47: 0x104d6fc in forward (test)
const result = try Tensor.from_ndarray(
^
/home/am/work/zignet/src/tensor.zig:326:30: 0x104537e in matmul (test)
return try op.forward(self.allocator, &[_]*Tensor{ self, other });
^
/home/am/work/zignet/src/layers.zig:80:47: 0x1044da7 in forward (test)
const matmul_result = try input.matmul(weights_transposed);
^
/home/am/work/zignet/src/layers.zig:16:47: 0x1045f0b in forward (test)
.Linear => |linear| linear.forward(input),
^
/home/am/work/zignet/src/sequential.zig:26:39: 0x1046033 in forward (test)
output = try layer.forward(output);
^
/home/am/work/zignet/src/sequential.zig:82:37: 0x10465a6 in test.Model and Layer functionality (test)
const output = try model.forward(input);
^
/home/am/zig/lib/compiler/test_runner.zig:208:25: 0x10f7bdc in mainTerminal (test)
if (test_fn.func()) |_| {
^
/home/am/zig/lib/compiler/test_runner.zig:57:28: 0x10f1d40 in main (test)
return mainTerminal();
^
/home/am/zig/lib/std/start.zig:647:22: 0x10f1296 in posixCallMainAndExit (test)
root.main();
^
All 1 tests passed.
25 errors were logged.
1 tests leaked memory.
It looks like your matmul_result is being allocated but you never deallocate it. I dont know know enough about your lifetimes to be sure, but it may be as simple as defering matmul_result.deinit() since it seems to be an intermediate value.
However you probaly dont want to allocate an intermediate value. It wpuld probably be beter to do an in place mutation or let it be stack allocated.
If you look two lines below, he reassigns the output var to a new pointer, so no douple free here unless there are no layers. However this will be a memory leak when there are more than 1 layers.
This was a good catch and helped reduced the amount of errors I have.
As @Southporter predicted, fixing the matmul stuff, made the one layer test pass but not the two layers
I use this instead now
pub fn forward(self: *Model, input: *Tensor) !*Tensor {
// Clone the input tensor to avoid modifying the original input
var output = try input.clone();
errdefer output.deinit(); // Ensure output is deinitialized if an error occurs
for (self.layers.items) |layer| {
const new_output = try layer.forward(output);
output.deinit(); // Deinitialize the old output tensor
output = new_output;
}
return output;
}
This seems to solve it. I still get a segmentation fault
Segmentation fault at address 0x7f28c8455110
/home/am/work/zignet/src/tensor.zig:63:33: 0x104622b in deinit (test)
self.allocator.free(self.shape);
^
/home/am/work/zignet/src/layers.zig:111:28: 0x1052ebb in deinit (test)
self.weights.deinit();
^
/home/am/work/zignet/src/layers.zig:34:46: 0x104cd47 in deinit (test)
.Linear => |linear| linear.deinit(),
^
/home/am/work/zignet/src/sequential.zig:58:25: 0x1048caa in deinit (test)
layer.deinit();
^
/home/am/work/zignet/src/sequential.zig:108:23: 0x10fabfa in test.Model and Layer functionality with 2 layers (test)
defer model.deinit();
^
/home/am/zig/lib/compiler/test_runner.zig:208:25: 0x110275c in mainTerminal (test)
if (test_fn.func()) |_| {
^
/home/am/zig/lib/compiler/test_runner.zig:57:28: 0x10fc8c0 in main (test)
return mainTerminal();
^
/home/am/zig/lib/std/start.zig:647:22: 0x10fbe16 in posixCallMainAndExit (test)
root.main();
^
/home/am/zig/lib/std/start.zig:271:5: 0x10fba4d in _start (test)
asm volatile (switch (native_arch) {
^
???:?:?: 0x0 in ??? (???)
error: the following test command crashed:
but at least the original issue seems solved
I think I solved this by avoiding the double free.
const input_size: usize = 2;
const hidden_size: usize = 6;
const output_size: usize = 2;
var linear_layer1 = try Linear.init(allocator, input_size, hidden_size, activations.relu);
var linear_layer2 = try Linear.init(allocator, hidden_size, output_size, activations.relu);
// Initialize a model
var model = try Model.init(allocator);
defer model.deinit();
// Wrap the Linear and ReLU layers in a Layer union
var linear_layer_wrapper1 = Layer{ .Linear = &linear_layer1 };
var linear_layer_wrapper2 = Layer{ .Linear = &linear_layer2 };
// Add the layers to the model
try model.add_layer(&linear_layer_wrapper1);
try model.add_layer(&linear_layer_wrapper2);
This way the layers are freed using the model only.
Thank you both for your help.
1 Like