Add check for malloc to prevent NULL pointer dereference in docs exam…

…ple. (#2922)


Signed-off-by: Sunny-Anand <[email protected]>

docs/doc_example/main.c

@@ -1,5 +1,6 @@
 #include <OnnxMlirRuntime.h>
 #include <stdio.h>
+#include <assert.h>
 
 OMTensorList *run_main_graph(OMTensorList *);
 
@@ -11,9 +12,16 @@ OMTensorList *create_input_list() {
 
   // Construct float arrays filled with 1s or 2s.
   float *x1Data = (float *)malloc(sizeof(float) * num_elements);
+  // Check if memory is allocated for generating the data.
+  if(!x1Data) return NULL;
   for (int i = 0; i < num_elements; i++)
     x1Data[i] = 1.0;
   float *x2Data = (float *)malloc(sizeof(float) * num_elements);
+  // Check if memory is allocated for generating the data.
+  if(!x2Data){
+    free(x1Data);
+    return NULL;
+  }
   for (int i = 0; i < num_elements; i++)
     x2Data[i] = 2.0;
 
@@ -32,7 +40,10 @@ OMTensorList *create_input_list() {
 int main() {
   // Generate input TensorList
   OMTensorList *input_list = create_input_list();
-
+  if(!input_list){
+    // Return 2 for failure to create inputs.
+    return 2;
+  }
   // Call the compiled onnx model function.
   OMTensorList *output_list = run_main_graph(input_list);
   if (!output_list) {