[Experimental] Enable kleidi AI examples to run on graviton3

Author: akote123

torchao/experimental/CMakeLists.txt

@@ -29,13 +29,17 @@ if(TORCHAO_BUILD_KLEIDIAI)
 endif()
 include(CMakePrintHelpers)
 
-add_compile_options("-Wall" "-Werror" "-Wno-deprecated")
+if(CMAKE_SYSTEM_PROCESSOR STREQUAL "aarch64") 
+    add_compile_options("-Wall" "-Werror" "-Wno-deprecated" "-march=armv8.2-a+dotprod" "-fPIC" "-Wno-error=unknown-pragmas") 
+else()
+    add_compile_options("-Wall" "-Werror" "-Wno-deprecated")   
+endif()
 
 include(CMakePrintHelpers)
 message("TORCHAO_INCLUDE_DIRS: ${TORCHAO_INCLUDE_DIRS}")
 include_directories(${TORCHAO_INCLUDE_DIRS})
 
-if(CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
+if(CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64" OR CMAKE_SYSTEM_PROCESSOR STREQUAL "aarch64")
     if(TORCHAO_BUILD_KLEIDIAI)
         message(STATUS "Building with Arm KleidiAI library")
         add_compile_definitions(TORCHAO_ENABLE_KLEIDI=1)

torchao/experimental/kernels/cpu/aarch64/valpacking/interleave.cpp

@@ -7,6 +7,7 @@
 #include <torchao/experimental/kernels/cpu/aarch64/valpacking/valpack.h>
 #include <cassert>
 #include <cstring>
+#include <cstdint>
 
 // Interleaves data across channels (row/column) and groups.
 // Each channel is the same size (vals_per_channel) and is

torchao/experimental/ops/packed_weights_header.h

@@ -51,7 +51,7 @@ class PackedWeightsHeader {
     auto header = reinterpret_cast<const int*>(packed_weights);
     assert(header[0] == PackedWeightsHeader::magic);
     params_type params;
-    for (int i = 0; i < params.size(); i++) {
+    for (size_t i = 0; i < params.size(); i++) {
       params[i] = header[i + 2];
     }
     return PackedWeightsHeader(
@@ -62,7 +62,7 @@ class PackedWeightsHeader {
     if (type != other.type) {
       return false;
     }
-    for (int i = 0; i < params.size(); i++) {
+    for (size_t i = 0; i < params.size(); i++) {
       if (params[i] != other.params[i]) {
         return false;
       }

torchao/experimental/ops/parallel-aten-impl.h

@@ -5,7 +5,7 @@
 // LICENSE file in the root directory of this source tree.
 
 #pragma once
-#include <Aten/Parallel.h>
+#include <torchao/experimental/ops/parallel.h>
 #include <torch/library.h>
 #include <torch/torch.h>