From 78ee555cefce3245dea6952b88dab1e1a5ab9020 Mon Sep 17 00:00:00 2001
From: Tom <tom@example.com>
Date: Thu, 29 Jan 2026 19:19:34 +0100
Subject: [PATCH] Add security and stability validation to generated BitNet
 kernels

---
 utils/codegen_tl1.py | 14 ++++++++++++++
 utils/codegen_tl2.py | 14 ++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/utils/codegen_tl1.py b/utils/codegen_tl1.py
index 4c2e7dd3f..36065a87f 100644
--- a/utils/codegen_tl1.py
+++ b/utils/codegen_tl1.py
@@ -345,12 +345,26 @@ def gen_transform_code(kernel_shape):
     }}\n".format(kernel_shapes[i][0], kernel_shapes[i][1])])
 
     kernel_code = "".join([kernel_code, "\n\
+    if (bm == 0) {{\n\
+        fprintf(stderr, \"ggml_bitnet_transform_tensor: unsupported tensor dimensions m=%d, k=%d\\n\", m, k);\n\
+        return;\n\
+    }}\n\
+\n\
+    if (tensor->data == nullptr) {{\n\
+        fprintf(stderr, \"ggml_bitnet_transform_tensor: tensor->data is null\\n\");\n\
+        return;\n\
+    }}\n\
+\n\
     const int n_tile_num = m / bm;\n\
     const int BK = bk;\n\
     uint8_t * qweights;\n\
     bitnet_float_type * scales;\n\
 \n\
     scales = (bitnet_float_type *) aligned_malloc(sizeof(bitnet_float_type));\n\
+    if (scales == nullptr) {{\n\
+        fprintf(stderr, \"ggml_bitnet_transform_tensor: failed to allocate scales\\n\");\n\
+        return;\n\
+    }}\n\
     qweights = (uint8_t *) tensor->data;\n\
     float * i2_scales = (float * )(qweights + k * m / 4);\n\
     scales[0] = (bitnet_float_type) i2_scales[0];\n\
diff --git a/utils/codegen_tl2.py b/utils/codegen_tl2.py
index 4d9408123..6c0a473bc 100644
--- a/utils/codegen_tl2.py
+++ b/utils/codegen_tl2.py
@@ -649,12 +649,26 @@ def gen_transform_code(kernel_shapes):
     }}\n".format(kernel_shapes[i][0], kernel_shapes[i][1])])
 
     kernel_code = "".join([kernel_code, "\n\
+    if (bm == 0) {{\n\
+        fprintf(stderr, \"ggml_bitnet_transform_tensor: unsupported tensor dimensions m=%d, k=%d\\n\", m, k);\n\
+        return;\n\
+    }}\n\
+\n\
+    if (tensor->data == nullptr) {{\n\
+        fprintf(stderr, \"ggml_bitnet_transform_tensor: tensor->data is null\\n\");\n\
+        return;\n\
+    }}\n\
+\n\
     const int n_tile_num = m / bm;\n\
     const int BK = bk;\n\
     uint8_t * qweights;\n\
     bitnet_float_type * scales;\n\
 \n\
     scales = (bitnet_float_type *) aligned_malloc(sizeof(bitnet_float_type));\n\
+    if (scales == nullptr) {{\n\
+        fprintf(stderr, \"ggml_bitnet_transform_tensor: failed to allocate scales\\n\");\n\
+        return;\n\
+    }}\n\
     qweights = (uint8_t *) tensor->data;\n\
     int nbytes = (k - 256) * m / 3 * 5 / 8 + 256 * m / 2 * 4 / 8;\n\
     if (nbytes % 32 != 0) nbytes = 32 - nbytes % 32 + nbytes;\n\