Skip to content

Commit 315cfbe

Browse files
akapoor3518akapoor3518
authored
Merge pull request #56 from tsisw/FIR-980
@FIR-980 - llama.cpp: RMS_NORM Kernel implementation
2 parents a1ffe42 + 5057ad6 commit 315cfbe

File tree

5 files changed

+162
-21
lines changed

5 files changed

+162
-21
lines changed

‎examples/simple/simple-backend-tsi.cpp‎

Lines changed: 61 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@ float test_input_1[GGML_TSAVORITE_KERNEL_TYPE_COUNT][NUM_ELEMENTS] = {
3939
{1.1, -4.4, 10, -5, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, -23, 24, 25, -26, 27, -28, 29, -30, 31, -32.6},
4040
//SIN Kernel
4141
{1.1, 4.4, 10, 5, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 20, 20, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32.6},
42+
//RMS_NORM Kernel
43+
{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32},
4244
//SIGMOID Kernel need to fix not tested
4345
{1.1, 4.4, 10, 5, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 20, 20, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32.6},
4446
//SILU Kernel
@@ -64,6 +66,8 @@ float test_input_2[GGML_TSAVORITE_KERNEL_TYPE_COUNT][NUM_ELEMENTS] = {
6466
{1.1, 2.2, 5, 10, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32},
6567
//SIN Kernel input not used
6668
{1.1, 2.2, 5, 10, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32},
69+
//RMS_NORM Kernel input is not used
70+
{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32},
6771
//SIGMOID Kernel not used
6872
{1.1, 4.4, 10, 5, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 20, 20, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32.6},
6973
//SILU Kernel not used
@@ -89,11 +93,13 @@ float test_result[GGML_TSAVORITE_KERNEL_TYPE_COUNT][NUM_ELEMENTS] = {
8993
{1.1, 4.4, 10, 5, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32.6},
9094
//SIN Kernel
9195
{0.891207, -0.951602, -0.544021, -0.958924, -0.958924, -0.279416, 0.656987, 0.989358, 0.412118, -0.544021, -0.999990, -0.536573, 0.420167, 0.990607, 0.650288, -0.287903, -0.961398, -0.750987, 0.149877, 0.912945, 0.912945, 0.912945, -0.846220, -0.905578, -0.132352, 0.762559, 0.956376, 0.270906, -0.663634, -0.988032, -0.404039, 0.926149},
96+
//RMS_NORM Kernel
97+
{0.052888, 0.105776, 0.158664, 0.211552, 0.264440, 0.317328, 0.370216, 0.423104, 0.475992, 0.528880, 0.581768, 0.634656, 0.687544, 0.740432, 0.793320, 0.846208, 0.899096, 0.951984, 1.004872, 1.057760, 1.110648, 1.163536, 1.216424, 1.269312, 1.322200, 1.375088, 1.427976, 1.480864, 1.533752, 1.586640, 1.639528, 1.692416},
9298
//SIGMOID Kernel not tested
9399
{0.891207, -0.951602, -0.544021, -0.958924, -0.958924, -0.279416, 0.656987, 0.989358, 0.412118, -0.544021, -0.999990, -0.536573, 0.420167, 0.990607, 0.650288, -0.287903, -0.961398, -0.750987, 0.149877, 0.912945, 0.912945, 0.912945, -0.846220, -0.905578, -0.132352, 0.762559, 0.956376, 0.270906, -0.663634, -0.988032, -0.404039, 0.926149},
94100
// SILU Kernel
95101
{-0.000002, -0.000005, -0.000012, -0.000029, -0.000074, -0.000184, -0.000454, -0.001111, -0.002683, -0.006377, -0.014836, -0.033464, -0.071945, -0.142278, -0.238406, -0.268941, 0.000000, 0.731059, 1.761594, 2.857722, 3.928055, 4.966536, 5.985164, 6.993623, 7.997317, 8.998889, 9.999546, 10.999816, 11.999926, 12.999971, 13.999988, 14.999995}
96-
102+
97103
};
98104

99105
float test_input_scale_1[GGML_TSAVORITE_KERNEL_TYPE_COUNT][NUM_ELEMENTS_SCALE] = {
@@ -151,6 +157,12 @@ float test_input_scale_1[GGML_TSAVORITE_KERNEL_TYPE_COUNT][NUM_ELEMENTS_SCALE] =
151157
-16, 25, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
152158
-1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
153159
-1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1},
160+
//RMS_NORM Kernel
161+
{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32,
162+
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32,
163+
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32,
164+
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32,
165+
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25},
154166
//SIGMOID KERNEL need to fix input data
155167
{-1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
156168
-9, 4, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
@@ -217,6 +229,12 @@ float test_input_scale_2[GGML_TSAVORITE_KERNEL_TYPE_COUNT][NUM_ELEMENTS_SCALE] =
217229
-16, 25, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
218230
-1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
219231
-1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1},
232+
//RMS_NORM Kernel input not used
233+
{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32,
234+
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32,
235+
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32,
236+
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32,
237+
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25},
220238
//SIGMOID KERNEL input not used
221239
{-1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
222240
-9, 4, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
@@ -291,6 +309,24 @@ float test_result_scale[GGML_TSAVORITE_KERNEL_TYPE_COUNT][NUM_ELEMENTS_SCALE] =
291309
-0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471,
292310
0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471,
293311
0.841471, 0.841471, 0.841471},
312+
//RMS_NORM Kernel
313+
{
314+
0.054620, 0.109240, 0.163860, 0.218479, 0.273099, 0.327719, 0.382339, 0.436959, 0.491579, 0.546199,
315+
0.600818, 0.655438, 0.710058, 0.764678, 0.819298, 0.873918, 0.928537, 0.983157, 1.037777, 1.092397,
316+
1.147017, 1.201637, 1.256257, 1.310876, 1.365496, 1.420116, 1.474736, 1.529356, 1.583976, 1.638596,
317+
1.693215, 1.747835, 0.054620, 0.109240, 0.163860, 0.218479, 0.273099, 0.327719, 0.382339, 0.436959,
318+
0.491579, 0.546199, 0.600818, 0.655438, 0.710058, 0.764678, 0.819298, 0.873918, 0.928537, 0.983157,
319+
1.037777, 1.092397, 1.147017, 1.201637, 1.256257, 1.310876, 1.365496, 1.420116, 1.474736, 1.529356,
320+
1.583976, 1.638596, 1.693215, 1.747835, 0.054620, 0.109240, 0.163860, 0.218479, 0.273099, 0.327719,
321+
0.382339, 0.436959, 0.491579, 0.546199, 0.600818, 0.655438, 0.710058, 0.764678, 0.819298, 0.873918,
322+
0.928537, 0.983157, 1.037777, 1.092397, 1.147017, 1.201637, 1.256257, 1.310876, 1.365496, 1.420116,
323+
1.474736, 1.529356, 1.583976, 1.638596, 1.693215, 1.747835, 0.054620, 0.109240, 0.163860, 0.218479,
324+
0.273099, 0.327719, 0.382339, 0.436959, 0.491579, 0.546199, 0.600818, 0.655438, 0.710058, 0.764678,
325+
0.819298, 0.873918, 0.928537, 0.983157, 1.037777, 1.092397, 1.147017, 1.201637, 1.256257, 1.310876,
326+
1.365496, 1.420116, 1.474736, 1.529356, 1.583976, 1.638596, 1.693215, 1.747835, 0.054620, 0.109240,
327+
0.163860, 0.218479, 0.273099, 0.327719, 0.382339, 0.436959, 0.491579, 0.546199, 0.600818, 0.655438,
328+
0.710058, 0.764678, 0.819298, 0.873918, 0.928537, 0.983157, 1.037777, 1.092397, 1.147017, 1.201637,
329+
1.256257, 1.310876, 1.365496},
294330
// SIGMOID KERNEL, result need to change
295331
{-0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471,
296332
0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471, 0.841471,
@@ -335,14 +371,15 @@ static void ggml_log_callback_default(ggml_log_level level, const char * text, v
335371
}
336372

337373

338-
// --- FLOAT COMPARATOR
374+
// --- FLOAT COMPARATOR
339375
static bool ggml_tsi_compare_two_float(float a, float b) {
340376
// For very small values, use absolute error
341377
if (fabsf(a) < 1e-2f && fabsf(b) < 1e-2f) {
342378
return fabsf(a - b) < 1e-6f; // Accept up to 1e-6 difference for small values
343379
}
344-
// For larger values, use relative error
345-
const float epsilon = 1e-4f;
380+
// For larger values, use relative error with increased tolerance
381+
// Increased to 1e-3 (0.1%) to handle floating-point precision differences
382+
const float epsilon = 1e-3f; // Changed from 1e-4f to 1e-3f
346383
float diff = fabsf(a - b);
347384
float max_val = fmaxf(fabsf(a), fabsf(b));
348385
return diff < epsilon * max_val;
@@ -376,7 +413,7 @@ static bool load_model(simple_model & model, float * a, float * b, enum ggml_typ
376413
/*.mem_buffer =*/ NULL,
377414
/*.no_alloc =*/ true,
378415
};
379-
fprintf(stderr, "\n Calculating mem_size %ld %d and creating ggml context \n", ggml_tensor_overhead(), num_tensors);
416+
fprintf(stderr, "\n Calculating mem_size %ld %d and creating ggml context \n", ggml_tensor_overhead(), num_tensors);
380417

381418
// create context
382419
model.ctx = ggml_init(params);
@@ -475,6 +512,9 @@ static struct ggml_cgraph * build_graph(const simple_model& model, enum ggml_tsa
475512
case GGML_TSAVORITE_KERNEL_TYPE_SIN:
476513
result = ggml_sin(ctx0, model.a);
477514
break;
515+
case GGML_TSAVORITE_KERNEL_TYPE_RMS_NORM:
516+
result = ggml_rms_norm(ctx0, model.a, 1e-5);
517+
break;
478518
case GGML_TSAVORITE_KERNEL_TYPE_SIGMOID:
479519
result = ggml_sigmoid(ctx0, model.a);
480520
break;
@@ -500,11 +540,11 @@ static struct ggml_tensor * compute(const simple_model & model, ggml_gallocr_t a
500540

501541
fprintf(stderr, "\n Under Test case for compute API creating build_graph \n");
502542
struct ggml_cgraph * gf = build_graph(model, ops_type);
503-
if (!gf) {
543+
if (!gf) {
504544
fprintf(stderr, "\ncompute failed\n");
505545
return NULL;
506546
}
507-
547+
508548
// allocate tensors
509549
ggml_gallocr_alloc_graph(allocr, gf);
510550

@@ -533,6 +573,8 @@ enum ggml_tsavorite_kernel_type convert_testcase_to_ops_type (const char *testCa
533573
return GGML_TSAVORITE_KERNEL_TYPE_ABS;
534574
else if (!strcmp(testCase,"sin"))
535575
return GGML_TSAVORITE_KERNEL_TYPE_SIN;
576+
else if (!strcmp(testCase,"rms_norm"))
577+
return GGML_TSAVORITE_KERNEL_TYPE_RMS_NORM;
536578
else if (!strcmp(testCase,"sigmoid"))
537579
return GGML_TSAVORITE_KERNEL_TYPE_SIGMOID;
538580
else if (!strcmp(testCase,"silu"))
@@ -561,7 +603,10 @@ const char* convert_ops_type_to_testcase(enum ggml_tsavorite_kernel_type ops_typ
561603
return "neg";
562604
case GGML_TSAVORITE_KERNEL_TYPE_ABS:
563605
return "abs";
564-
case GGML_TSAVORITE_KERNEL_TYPE_SIN:
606+
case GGML_TSAVORITE_KERNEL_TYPE_SIN:
607+
return "sin";
608+
case GGML_TSAVORITE_KERNEL_TYPE_RMS_NORM:
609+
return "rms_norm";
565610
return "sin";
566611
case GGML_TSAVORITE_KERNEL_TYPE_SIGMOID:
567612
return "sigmoid";
@@ -601,26 +646,27 @@ int main(int argc, char *argv[]) {
601646
ops_type == GGML_TSAVORITE_KERNEL_TYPE_NEG ||
602647
ops_type == GGML_TSAVORITE_KERNEL_TYPE_ABS ||
603648
ops_type == GGML_TSAVORITE_KERNEL_TYPE_SIN ||
649+
ops_type == GGML_TSAVORITE_KERNEL_TYPE_RMS_NORM ||
604650
ops_type == GGML_TSAVORITE_KERNEL_TYPE_SIGMOID ||
605651
ops_type == GGML_TSAVORITE_KERNEL_TYPE_SILU)
606652
num_of_input_tensors = NUM_INPUT_URINARY_TENSORS;
607-
else
653+
else
608654
num_of_input_tensors = NUM_INPUT_TENSORS;
609655

610656
if (data_scale) {
611657
input1[ops_type] = test_input_scale_1[ops_type];
612-
elements_A = NUM_ELEMENTS_SCALE;
658+
elements_A = NUM_ELEMENTS_SCALE;
613659
if (num_of_input_tensors != NUM_INPUT_URINARY_TENSORS) {
614660
input2[ops_type] = test_input_scale_2[ops_type];
615-
elements_B = NUM_ELEMENTS_SCALE;
661+
elements_B = NUM_ELEMENTS_SCALE;
616662
}
617663
result_data[ops_type] = test_result_scale[ops_type];
618664
} else {
619665
input1[ops_type] = test_input_1[ops_type];
620-
elements_A = NUM_ELEMENTS;
666+
elements_A = NUM_ELEMENTS;
621667
if (num_of_input_tensors != NUM_INPUT_URINARY_TENSORS) {
622668
input2[ops_type] = test_input_2[ops_type];
623-
elements_B = NUM_ELEMENTS;
669+
elements_B = NUM_ELEMENTS;
624670
}
625671
result_data[ops_type] = test_result[ops_type];
626672
}
@@ -687,6 +733,8 @@ int main(int argc, char *argv[]) {
687733

688734
if (test_case_flag == false) {
689735
fprintf(stderr, "\n\n TEST CASE FAILED \n\n");
736+
ggml_free(model.ctx);
737+
ggml_backend_free(model.backend);
690738
return -1;
691739
}
692740
fprintf(stderr, "\n\n TEST CASE PASSED \n\n");

‎ggml-tsi-kernel‎

‎ggml/include/ggml-tsavorite.h‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,7 @@ enum ggml_tsavorite_kernel_type {
126126
GGML_TSAVORITE_KERNEL_TYPE_NEG,
127127
GGML_TSAVORITE_KERNEL_TYPE_ABS,
128128
GGML_TSAVORITE_KERNEL_TYPE_SIN,
129+
GGML_TSAVORITE_KERNEL_TYPE_RMS_NORM,
129130
GGML_TSAVORITE_KERNEL_TYPE_SIGMOID,
130131
GGML_TSAVORITE_KERNEL_TYPE_SILU,
131132

@@ -162,11 +163,15 @@ extern void _mlir_ciface_txe_abs_host(void *a, void *res);
162163
extern void _mlir_ciface_txe_sin_host(void *a, void *res);
163164
extern void _mlir_ciface_txe_sigmoid_host(void *a, void *res);
164165
extern void _mlir_ciface_txe_silu_host(void *a, void *res);
166+
extern void _mlir_ciface_txe_rms_norm_host(void *a, void *res, void *buf);
167+
165168
extern void ggml_tsi_log_tensor_data(tensor_log log_data);
166169

167170
#define NUM_OF_TXES 1
168-
// GML supports a maximum tensor rank of 4
171+
172+
// GGML supports tensors with a maximum rank of 4
169173
#define MEM_REF_DESCRIPTOR_RANK 4
174+
#define TSI_TVU_LOAD_SIZE 32
170175

171176
//
172177
// backend API

‎ggml/src/ggml-tsavorite/ggml-tsavorite.cpp‎

Lines changed: 90 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -458,6 +458,11 @@ static txe_compute_pipeline_state_s tsi_kernel_setup(enum ggml_tsavorite_kernel_
458458
kernel_pipeline->kernel_name = "TXE_SILU";
459459
flag = true;
460460
break;
461+
case GGML_TSAVORITE_KERNEL_TYPE_RMS_NORM:
462+
kernel_pipeline->_mlir_fptr_2_input = &_mlir_ciface_txe_rms_norm_host;
463+
kernel_pipeline->kernel_name = "TXE_RMS_NORM";
464+
flag = true;
465+
break;
461466
default:
462467
break;
463468
}
@@ -605,6 +610,7 @@ static struct ggml_backend_tsavorite_context *ggml_tsavorite_init(ggml_backend_d
605610
GGML_TSAVORITE_KERNEL(GGML_TSAVORITE_KERNEL_TYPE_SIN, true);
606611
GGML_TSAVORITE_KERNEL(GGML_TSAVORITE_KERNEL_TYPE_SIGMOID, true);
607612
GGML_TSAVORITE_KERNEL(GGML_TSAVORITE_KERNEL_TYPE_SILU, true);
613+
GGML_TSAVORITE_KERNEL(GGML_TSAVORITE_KERNEL_TYPE_RMS_NORM, true);
608614
}
609615

610616
GGML_TSAVORITE_LOG_INFO("End %s\n", __func__);
@@ -708,6 +714,8 @@ static bool ggml_tsavorite_supports_op(const struct ggml_backend_tsavorite_devic
708714
case GGML_OP_SQR:
709715
case GGML_OP_SIN:
710716
break;
717+
case GGML_OP_RMS_NORM:
718+
break;
711719
case GGML_OP_UNARY:
712720
switch (ggml_get_unary_op(op)) {
713721
case GGML_UNARY_OP_NEG:
@@ -755,6 +763,32 @@ static void ggml_tsavorite_decompose_unary_kernel(uint32_t num_elem, ggml_tensor
755763
return;
756764
}
757765

766+
template<int Rank>
767+
// Assumes tsi_alloc is available and returns a pointer to allocated memory
768+
static MemRefDescriptor<Rank>* create_mlir_buf(int K) {
769+
// TVU load size (e.g., 32 for 1024-bit vector with 32-bit elements)
770+
const int32_t tvu_size = TSI_TVU_LOAD_SIZE;
771+
772+
// Round up K to the next multiple of tvu_size
773+
int32_t num_of_elem = ((K % tvu_size) != 0) ? ((K / tvu_size) + 1) * tvu_size : K;
774+
775+
// Allocate memory dynamically: space for header + data
776+
MemRefDescriptor<Rank>* header = (MemRefDescriptor<Rank>*) tsi_alloc(
777+
sizeof(MemRefDescriptor<Rank>) + num_of_elem * sizeof(float)
778+
);
779+
780+
if (!header) {
781+
return header;
782+
}
783+
// Advance pointer to skip header and get to data
784+
int32_t* data = (int32_t*)(header + 1);
785+
786+
for (int32_t i = 0; i < num_of_elem; ++i) {
787+
data[i] = 0;
788+
}
789+
return header;
790+
}
791+
758792
// nodes are intermediate which has multiple src tensors & operation
759793
// Here we create multiple thread
760794
// Each Thread run the command buffer & pick Tensor and execute and get the result back base on
@@ -864,6 +898,10 @@ static enum ggml_status ggml_tsavorite_graph_compute(ggml_backend_t backend,
864898
kernel_type = GGML_TSAVORITE_KERNEL_TYPE_SIN;
865899
num_of_input_tensors = TSAVORITE_UNARY_INPUT_TENSORS;
866900
break;
901+
case GGML_OP_RMS_NORM:
902+
kernel_type = GGML_TSAVORITE_KERNEL_TYPE_RMS_NORM;
903+
num_of_input_tensors = TSAVORITE_UNARY_INPUT_TENSORS;
904+
break;
867905
case GGML_OP_UNARY:
868906
switch (ggml_get_unary_op(node)) {
869907
case GGML_UNARY_OP_NEG:
@@ -1079,8 +1117,54 @@ static enum ggml_status ggml_tsavorite_graph_compute(ggml_backend_t backend,
10791117
nodeP->shape[0] = num_elem_src0;
10801118
srcP0->strides[0] = 0;
10811119
nodeP->strides[0] = 0;
1082-
// kernel call
1083-
ctx->kernels[kernel_type].pipeline->_mlir_fptr_1_input(srcP0, nodeP);
1120+
1121+
if (kernel_type == GGML_TSAVORITE_KERNEL_TYPE_RMS_NORM) {
1122+
// tsi_alloc is invoked within the function below.
1123+
// We allocate 64 elements for RMS normalization used in the RMS kernel.
1124+
// Although only 32 elements are strictly necessary, reducing this would require changes to the RMS kernel.
1125+
// The remaining 32 elements are used to store src0->ne[0], replicated across each of the last 32 entries.
1126+
1127+
MemRefDescriptor<Rank>* buf = create_mlir_buf<Rank>(96);
1128+
1129+
if (!buf) {
1130+
GGML_TSAVORITE_LOG_ERROR("tsi_alloc failied for creating memory for buf \n");
1131+
return GGML_STATUS_ABORTED;
1132+
}
1133+
buf->offset = 0;
1134+
buf->data = buf->base = (void *)(buf+1);
1135+
1136+
float *val = (float *)buf->data;
1137+
int i;
1138+
for(i=64; i <= 95; ++i)
1139+
val[i] = node->ne[0];
1140+
1141+
int max_dim_index = GGML_MAX_DIMS -1;
1142+
int strides = 1;
1143+
bool flag = true;
1144+
for ( i = 0; i <= max_dim_index && src0->nb[i] != 0; ++i) {
1145+
if (src0->ne[i] == 0) {
1146+
srcP0->shape[max_dim_index - i] = 1;
1147+
nodeP->shape[max_dim_index - i] = 1;
1148+
flag = false;
1149+
}
1150+
else {
1151+
srcP0->shape[max_dim_index - i] = src0->ne[i];
1152+
nodeP->shape[max_dim_index - i] = node->ne[i];
1153+
}
1154+
srcP0->strides[max_dim_index - i] = strides;
1155+
nodeP->strides[max_dim_index - i] = strides;
1156+
1157+
// avoiding the case when src0->ne[i] is zero
1158+
if (flag)
1159+
strides = strides * src0->ne[i];
1160+
}
1161+
1162+
ctx->kernels[kernel_type].pipeline->_mlir_fptr_2_input(srcP0, nodeP, buf);
1163+
}
1164+
else {
1165+
// kernel call
1166+
ctx->kernels[kernel_type].pipeline->_mlir_fptr_1_input(srcP0, nodeP);
1167+
}
10841168
++device->stats.op_run_count[kernel_type].num_of_kernel_call;
10851169

10861170
if (ggml_tsavorite_log_type_val == GGML_TSAVORITE_LOG_DEBUG) {
@@ -1380,7 +1464,9 @@ static size_t ggml_backend_tsavorite_buffer_type_get_alloc_size(ggml_backend_buf
13801464
"\n\n\n\n Calculating---- Alloc ----Size header %lu and data %lu \n\n\n\n ",
13811465
sizeof(tensor_data_header), ggml_nbytes(tensor));
13821466

1383-
return (sizeof(tensor_data_header) + ggml_nbytes(tensor));
1467+
// Add 128-byte buffer to avoid crossing memory boundaries during TVU 1024-bit operations.
1468+
// TVU processes data in 1024-bit chunks, so the last elements may exceed allocated space without this padding.
1469+
return (sizeof(tensor_data_header) + ggml_nbytes(tensor) + 128);
13841470

13851471
TSI_UNUSED(buft);
13861472
}
@@ -1803,6 +1889,7 @@ static bool ggml_backend_tsavorite_device_offload_op(ggml_backend_dev_t dev,
18031889
case GGML_OP_SQRT:
18041890
case GGML_OP_SQR:
18051891
case GGML_OP_SIN:
1892+
case GGML_OP_RMS_NORM:
18061893
break;
18071894
case GGML_OP_UNARY:
18081895
switch (ggml_get_unary_op(op)) {

0 commit comments

Comments
 (0)