Enable compilation tests for microbenchmarks (#8904)

This creates a microbenchmark suite and the adds the compilation to CI bots.

CMakeLists.txt

@@ -39,6 +39,7 @@ option(IREE_ENABLE_THREADING "Builds IREE in with thread library support." ON)
 option(IREE_BUILD_COMPILER "Builds the IREE compiler." ON)
 option(IREE_BUILD_TESTS "Builds IREE unit tests." ON)
 option(IREE_BUILD_BENCHMARKS "Builds IREE benchmark suites." OFF)
+option(IREE_BUILD_MICROBENCHMARKS "Builds IREE microbenchmark suites." OFF)
 option(IREE_BUILD_DOCS "Builds IREE docs." OFF)
 option(IREE_BUILD_SAMPLES "Builds IREE sample projects." ON)
 option(IREE_BUILD_TRACY "Builds tracy server tools." OFF)
@@ -239,6 +240,7 @@ include(iree_trace_runner_test)
 include(iree_native_test)
 include(iree_cc_binary_benchmark)
 include(iree_benchmark_suite)
+include(iree_microbenchmark_suite)
 include(iree_hal_cts_test_suite)
 
 set(DEFAULT_CMAKE_BUILD_TYPE "Release")
@@ -574,6 +576,11 @@ if(${IREE_BUILD_BENCHMARKS})
   add_custom_target(iree-benchmark-suites)
 endif()
 
+if(${IREE_BUILD_MICROBENCHMARKS})
+  # Add top-level custom targets to drive generating microbenchmark suites.
+  add_custom_target(iree-microbenchmark-suites)
+endif()
+
 if(${IREE_BUILD_DOCS})
   # Add a top-level custom target to drive generating all documentation.
   # Register it to the default target given that IREE_BUILD_DOCS is explicitly

build_tools/cmake/build_android_benchmark.sh

@@ -64,11 +64,13 @@ cd build-host
   -DIREE_BUILD_COMPILER=ON \
   -DIREE_BUILD_TESTS=OFF \
   -DIREE_BUILD_BENCHMARKS=ON \
+  -DIREE_BUILD_MICROBENCHMARKS=ON \
   -DIREE_BUILD_SAMPLES=OFF
 
 "${CMAKE_BIN}" --build . --target install -- -k 0
 # Also generate artifacts for benchmarking on Android.
 "${CMAKE_BIN}" --build . --target iree-benchmark-suites -- -k 0
+"${CMAKE_BIN}" --build . --target iree-microbenchmark-suites -- -k 0
 # --------------------------------------------------------------------------- #
 
 # --------------------------------------------------------------------------- #

build_tools/cmake/iree_microbenchmark_suite.cmake

@@ -0,0 +1,61 @@
+# Copyright 2022 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+# iree_microbenchmark_suite()
+#
+# Generates microbenchmark suites for MLIR input modules.
+# Parameters:
+#   NAME: Name of target.
+#   SRCS: Source files to compile into a bytecode module (list of strings).
+#   FLAGS: Flags to pass to the translation tool (list of strings).
+
+function(iree_microbenchmark_suite)
+  if(NOT IREE_BUILD_MICROBENCHMARKS)
+    return()
+  endif()
+
+  cmake_parse_arguments(
+    _RULE
+    ""
+    "NAME"
+    "SRCS;FLAGS"
+    ${ARGN}
+  )
+
+  iree_package_name(PACKAGE_NAME)
+
+  foreach(_SRC IN LISTS _RULE_SRCS)
+    set(_TRANSLATE_TOOL "iree-compile")
+    set(_TRANSLATE_SRC "${_SRC}")
+    set(_MODULE_FILE_NAME "${_RULE_NAME}_${_SRC}.vmfb")
+    set(_TARGET_NAME "${PACKAGE_NAME}_${_MODULE_FILE_NAME}")
+    iree_get_executable_path(_TRANSLATE_TOOL_EXECUTABLE "${_TRANSLATE_TOOL}")
+    set(_ARGS "${_RULE_FLAGS}")
+    get_filename_component(_TRANSLATE_SRC_PATH "${_TRANSLATE_SRC}" REALPATH)
+    list(APPEND _ARGS "${_TRANSLATE_SRC_PATH}")
+    list(APPEND _ARGS "-o")
+    list(APPEND _ARGS "${_MODULE_FILE_NAME}")
+
+    add_custom_command(
+      OUTPUT
+        "${_MODULE_FILE_NAME}"
+      COMMAND
+        "${_TRANSLATE_TOOL_EXECUTABLE}"
+        ${_ARGS}
+      # Changes to either the translation tool or the input source should
+      # trigger rebuilding.
+      DEPENDS
+        "${_TRANSLATE_TOOL_EXECUTABLE}"
+        "${_TRANSLATE_SRC}"
+      VERBATIM
+    )
+    add_custom_target("${_TARGET_NAME}"
+      DEPENDS
+        "${_MODULE_FILE_NAME}"
+    )
+    add_dependencies(iree-microbenchmark-suites "${_TARGET_NAME}")
+  endforeach(_SRC IN LISTS _SRCS)
+endfunction(iree_microbenchmark_suite)

build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-swiftshader-asan/build.sh

@@ -37,6 +37,9 @@ CMAKE_ARGS=(
   "-DIREE_ENABLE_ASAN=ON"
   "-B" "${CMAKE_BUILD_DIR?}"
 
+  # Also check if microbenchmarks are buildable.
+  "-DIREE_BUILD_MICROBENCHMARKS=ON"
+
   # Enable CUDA compiler and runtime builds unconditionally. Our CI images all
   # have enough deps to at least build CUDA support and compile CUDA binaries
   # (but not necessarily test on real hardware).
@@ -55,6 +58,10 @@ echo "Building test deps"
 echo "------------------"
 "${CMAKE_BIN?}" --build "${CMAKE_BUILD_DIR?}" --target iree-test-deps -- -k 0
 
+echo "Building microbenchmark suites"
+echo "------------------"
+"${CMAKE_BIN?}" --build "${CMAKE_BUILD_DIR?}" --target iree-microbenchmark-suites -- -k 0
+
 # Respect the user setting, but default to as many jobs as we have cores.
 export CTEST_PARALLEL_LEVEL=${CTEST_PARALLEL_LEVEL:-$(nproc)}
 

iree/test/microbenchmarks/CMakeLists.txt

@@ -1,24 +1,19 @@
-################################################################################
-# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from           #
-# iree/test/microbenchmarks/BUILD                                              #
-#                                                                              #
-# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary   #
-# CMake-only content.                                                          #
-#                                                                              #
-# To disable autogeneration for this file entirely, delete this header.        #
-################################################################################
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
 
 iree_add_all_subdirs()
 
-iree_lit_test_suite(
+iree_microbenchmark_suite(
   NAME
-    lit
+   "microbenchmark"
   SRCS
+    "linalg_mmt4d.mlir"
     "linalg_transpose.mlir"
-  TOOLS
-    iree::tools::iree-run-mlir
-  LABELS
-    "hostonly"
+    "mhlo_conv.mlir"
+    "mhlo_dot.mlir"
+    "mhlo_dot_general.mlir"
+    "mhlo_fft_abs.mlir"
+  FLAGS
+    "--iree-hal-target-backends=dylib-llvm-aot"
+    "--iree-input-type=mhlo"
+    "--iree-llvm-target-cpu-features=host"
 )
-
-### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###

iree/test/microbenchmarks/mhlo_conv.mlir

@@ -13,17 +13,17 @@ func.func @conv_244_112_3x3_3x32() -> tensor<1x112x112x32xf32> {
     %filter = util.unfoldable_constant dense<1.0> : tensor<3x3x3x32xf32>
     %0 = "mhlo.convolution"(%input, %filter) {
         batch_group_count = 1 : i64,
-        dimension_numbers = {
-            input_batch_dimension = 0 : i64,
-            input_feature_dimension = 3 : i64,
-            input_spatial_dimensions = dense<[1, 2]> : tensor<2xi64>,
-            kernel_input_feature_dimension = 2 : i64,
-            kernel_output_feature_dimension = 3 : i64,
-            kernel_spatial_dimensions = dense<[0, 1]> : tensor<2xi64>,
-            output_batch_dimension = 0 : i64,
-            output_feature_dimension = 3 : i64,
-            output_spatial_dimensions = dense<[1, 2]> : tensor<2xi64>
-        },
+        dimension_numbers = #mhlo.conv<raw
+            input_batch_dimension = 0,
+            input_feature_dimension = 3,
+            input_spatial_dimensions = [1, 2],
+            kernel_input_feature_dimension = 2,
+            kernel_output_feature_dimension = 3,
+            kernel_spatial_dimensions = [0, 1],
+            output_batch_dimension = 0,
+            output_feature_dimension = 3,
+            output_spatial_dimensions = [1, 2]
+        >,
         feature_group_count = 1 : i64,
         padding = dense<[[0, 1], [0, 1]]> : tensor<2x2xi64>,
         rhs_dilation = dense<1> : tensor<2xi64>,
@@ -37,17 +37,17 @@ func.func @conv_112_112_1x1_32x64() -> tensor<1x112x112x64xf32> {
     %filter = util.unfoldable_constant dense<1.0> : tensor<1x1x32x64xf32>
     %0 = "mhlo.convolution"(%input, %filter) {
         batch_group_count = 1 : i64,
-        dimension_numbers = {
-            input_batch_dimension = 0 : i64,
-            input_feature_dimension = 3 : i64,
-            input_spatial_dimensions = dense<[1, 2]> : tensor<2xi64>,
-            kernel_input_feature_dimension = 2 : i64,
-            kernel_output_feature_dimension = 3 : i64,
-            kernel_spatial_dimensions = dense<[0, 1]> : tensor<2xi64>,
-            output_batch_dimension = 0 : i64,
-            output_feature_dimension = 3 : i64,
-            output_spatial_dimensions = dense<[1, 2]> : tensor<2xi64>
-        },
+        dimension_numbers = #mhlo.conv<raw
+            input_batch_dimension = 0,
+            input_feature_dimension = 3,
+            input_spatial_dimensions = [1, 2],
+            kernel_input_feature_dimension = 2,
+            kernel_output_feature_dimension = 3,
+            kernel_spatial_dimensions = [0, 1],
+            output_batch_dimension = 0,
+            output_feature_dimension = 3,
+            output_spatial_dimensions = [1, 2]
+        >,
         feature_group_count = 1 : i64,
         padding = dense<0> : tensor<2x2xi64>,
         rhs_dilation = dense<1> : tensor<2xi64>,
@@ -61,17 +61,17 @@ func.func @conv_7_7_1x1_1024x1024() -> tensor<1x7x7x1024xf32> {
     %filter = util.unfoldable_constant dense<1.0> : tensor<1x1x1024x1024xf32>
     %0 = "mhlo.convolution"(%input, %filter) {
         batch_group_count = 1 : i64,
-        dimension_numbers = {
-            input_batch_dimension = 0 : i64,
-            input_feature_dimension = 3 : i64,
-            input_spatial_dimensions = dense<[1, 2]> : tensor<2xi64>,
-            kernel_input_feature_dimension = 2 : i64,
-            kernel_output_feature_dimension = 3 : i64,
-            kernel_spatial_dimensions = dense<[0, 1]> : tensor<2xi64>,
-            output_batch_dimension = 0 : i64,
-            output_feature_dimension = 3 : i64,
-            output_spatial_dimensions = dense<[1, 2]> : tensor<2xi64>
-        },
+        dimension_numbers = #mhlo.conv<raw
+            input_batch_dimension = 0,
+            input_feature_dimension = 3,
+            input_spatial_dimensions = [1, 2],
+            kernel_input_feature_dimension = 2,
+            kernel_output_feature_dimension = 3,
+            kernel_spatial_dimensions = [0, 1],
+            output_batch_dimension = 0,
+            output_feature_dimension = 3,
+            output_spatial_dimensions = [1, 2]
+        >,
         feature_group_count = 1 : i64,
         padding = dense<0> : tensor<2x2xi64>,
         rhs_dilation = dense<1> : tensor<2xi64>,
@@ -96,17 +96,17 @@ func.func @depthwise_conv_15x1_1x1_15x1_1x1024_1024() -> tensor<1x1x1x1024xf32>
   %filter = util.unfoldable_constant dense<1.0> : tensor<15x1x1x1024xf32>
   %res = "mhlo.convolution"(%input, %filter) {
     batch_group_count = 1 : i64,
-    dimension_numbers = {
-      input_batch_dimension = 0 : i64,
-      input_feature_dimension = 3 : i64,
-      input_spatial_dimensions = dense<[1, 2]> : tensor<2xi64>,
-      kernel_input_feature_dimension = 2 : i64,
-      kernel_output_feature_dimension = 3 : i64,
-      kernel_spatial_dimensions = dense<[0, 1]> : tensor<2xi64>,
-      output_batch_dimension = 0 : i64,
-      output_feature_dimension = 3 : i64,
-      output_spatial_dimensions = dense<[1, 2]> : tensor<2xi64>
-      },
+    dimension_numbers = #mhlo.conv<raw
+      input_batch_dimension = 0,
+      input_feature_dimension = 3,
+      input_spatial_dimensions = [1, 2],
+      kernel_input_feature_dimension = 2,
+      kernel_output_feature_dimension = 3,
+      kernel_spatial_dimensions = [0, 1],
+      output_batch_dimension = 0,
+      output_feature_dimension = 3,
+      output_spatial_dimensions = [1, 2]
+      >,
     feature_group_count = 1024 : i64,
     padding = dense<0> : tensor<2x2xi64>,
     rhs_dilation = dense<1> : tensor<2xi64>,
@@ -120,17 +120,17 @@ func.func @depthwise_conv_15x1_1x1_15x1_1x512_512() -> tensor<1x1x1x512xf32> {
   %filter = util.unfoldable_constant dense<1.0> : tensor<15x1x1x512xf32>
   %res = "mhlo.convolution"(%input, %filter) {
     batch_group_count = 1 : i64,
-    dimension_numbers = {
-      input_batch_dimension = 0 : i64,
-      input_feature_dimension = 3 : i64,
-      input_spatial_dimensions = dense<[1, 2]> : tensor<2xi64>,
-      kernel_input_feature_dimension = 2 : i64,
-      kernel_output_feature_dimension = 3 : i64,
-      kernel_spatial_dimensions = dense<[0, 1]> : tensor<2xi64>,
-      output_batch_dimension = 0 : i64,
-      output_feature_dimension = 3 : i64,
-      output_spatial_dimensions = dense<[1, 2]> : tensor<2xi64>
-      },
+    dimension_numbers = #mhlo.conv<raw
+      input_batch_dimension = 0,
+      input_feature_dimension = 3,
+      input_spatial_dimensions = [1, 2],
+      kernel_input_feature_dimension = 2,
+      kernel_output_feature_dimension = 3,
+      kernel_spatial_dimensions = [0, 1],
+      output_batch_dimension = 0,
+      output_feature_dimension = 3,
+      output_spatial_dimensions = [1, 2]
+    >,
     feature_group_count = 512 : i64,
     padding = dense<0> : tensor<2x2xi64>,
     rhs_dilation = dense<1> : tensor<2xi64>,
@@ -144,17 +144,17 @@ func.func @depthwise_conv_16x1_2x1_16x1_1x512_512() -> tensor<1x2x1x512xf32> {
   %filter = util.unfoldable_constant dense<1.0> : tensor<15x1x1x512xf32>
   %res = "mhlo.convolution"(%input, %filter) {
     batch_group_count = 1 : i64,
-    dimension_numbers = {
-      input_batch_dimension = 0 : i64,
-      input_feature_dimension = 3 : i64,
-      input_spatial_dimensions = dense<[1, 2]> : tensor<2xi64>,
-      kernel_input_feature_dimension = 2 : i64,
-      kernel_output_feature_dimension = 3 : i64,
-      kernel_spatial_dimensions = dense<[0, 1]> : tensor<2xi64>,
-      output_batch_dimension = 0 : i64,
-      output_feature_dimension = 3 : i64,
-      output_spatial_dimensions = dense<[1, 2]> : tensor<2xi64>
-    },
+    dimension_numbers = #mhlo.conv<raw
+      input_batch_dimension = 0,
+      input_feature_dimension = 3,
+      input_spatial_dimensions = [1, 2],
+      kernel_input_feature_dimension = 2,
+      kernel_output_feature_dimension = 3,
+      kernel_spatial_dimensions = [0, 1],
+      output_batch_dimension = 0,
+      output_feature_dimension = 3,
+      output_spatial_dimensions = [1, 2]
+    >,
     feature_group_count = 512 : i64,
     padding = dense<0> : tensor<2x2xi64>,
     rhs_dilation = dense<1> : tensor<2xi64>,