[Testing and Infra] LIT testing and passes enabling (#33)

Devjiu · web-flow · commit 294c4bd61c13 · 2024-05-07T16:43:32.000+02:00
This commit enables LIT tests and introduces example of MLP LIT test.
To run all of them use:
```sh
cmake --build . --target gc-check
```

Also introduced infrastructure for adding Passes and registered dummy TileNamedLinalg pass.
Currently it simply print matmuls. To enable use option: "--tile-named-linalg".

Signed-off-by: Dmitrii Makarenko &lt;dmitrii.makarenko@intel.com&gt;
diff --git a/include/gc-dialects/CMakeLists.txt b/include/gc-dialects/CMakeLists.txt
@@ -1,3 +1,8 @@
 add_subdirectory(OnednnGraph)
 add_subdirectory(Microkernel)
 add_subdirectory(Linalgx)
+
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name GraphCompiler)
+add_public_tablegen_target(GraphCompilerPassIncGen)
+add_mlir_doc(Passes GraphCompilerPasses ./ -gen-pass-doc)
diff --git a/include/gc-dialects/Passes.h b/include/gc-dialects/Passes.h
@@ -0,0 +1,25 @@
+//===- Passes.h - Graph Compiler passes -------------------------*- C++ -*-===//
+//
+// This file is 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef GC_PASSES_H
+#define GC_PASSES_H
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace gc {
+
+#define GEN_PASS_DECL
+#include "gc-dialects/Passes.h.inc"
+
+#define GEN_PASS_REGISTRATION
+#include "gc-dialects/Passes.h.inc"
+} // namespace gc
+} // namespace mlir
+
+#endif // GC_PASSES_H
diff --git a/include/gc-dialects/Passes.td b/include/gc-dialects/Passes.td
@@ -0,0 +1,20 @@
+//===- Passes.td - Graph Compiler passes -------------------*- tablegen -*-===//
+//
+// This file is 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef GC_DIALECT_GC_PASSES
+#define GC_DIALECT_GC_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def TileLinalgNamed : Pass<"tile-named-linalg", "func::FuncOp"> {
+  let summary = "Tile linalg named operations.";
+  let dependentDialects =
+      ["linalg::LinalgDialect", "scf::SCFDialect", "tensor::TensorDialect"];
+}
+
+#endif // GC_DIALECT_GC_PASSES
diff --git a/lib/gc-dialects/CMakeLists.txt b/lib/gc-dialects/CMakeLists.txt
@@ -1,3 +1,5 @@
 add_subdirectory(Linalgx)
 add_subdirectory(Microkernel)
 add_subdirectory(OnednnGraph)
+
+add_subdirectory(Transforms)
diff --git a/lib/gc-dialects/Transforms/CMakeLists.txt b/lib/gc-dialects/Transforms/CMakeLists.txt
@@ -0,0 +1,16 @@
+add_mlir_library(GCPasses
+  TileNamed.cpp
+
+  ADDITIONAL_HEADER_DIRS
+    ${PROJECT_SOURCE_DIR}/include/gc-dialects
+
+  DEPENDS
+    GraphCompilerPassIncGen
+
+  LINK_LIBS PUBLIC
+    ${mlir_dialect_libs}
+    MLIRIR
+    MLIRSupport
+    MLIRBufferizationToMemRef
+    MLIRBufferizationPipelines
+  )
diff --git a/lib/gc-dialects/Transforms/TileNamed.cpp b/lib/gc-dialects/Transforms/TileNamed.cpp
@@ -0,0 +1,49 @@
+//===- TileNamed.cpp - Tile Named Linalg Ops --------------------*- C++ -*-===//
+//
+// This file is 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "gc-dialects/Passes.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/SCF/Transforms/TileUsingInterface.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+
+namespace mlir {
+namespace gc {
+#define GEN_PASS_DEF_TILELINALGNAMED
+#include "gc-dialects/Passes.h.inc"
+} // namespace gc
+} // namespace mlir
+
+namespace {
+class TileLinalg : public mlir::gc::impl::TileLinalgNamedBase<TileLinalg> {
+
+  void runOnOperation() override {
+    auto *ctx = &getContext();
+    IRRewriter rewriter(ctx);
+
+    llvm::SmallVector<Operation *> to_tile;
+    for (Operation &o : getOperation()->getRegion(0).front().getOperations()) {
+      if (isa<linalg::MatmulOp>(o)) {
+        to_tile.push_back(&o);
+      }
+    }
+
+    for (Operation *o : to_tile) {
+      llvm::errs() << "func op body to tile: " << *o << "\n";
+    }
+  }
+};
+
+} // namespace
diff --git a/scripts/compile.sh b/scripts/compile.sh
@@ -21,5 +21,8 @@ get_llvm() (
 
 test -f "$llvm_dir/llvm-$llvm_hash"/llvm.tgz || get_llvm
 
-MLIR_DIR="$llvm_dir/lib/cmake/mlir" cmake -S . -G Ninja -B build
+cmake -S . -G Ninja -B build \
+    -DCMAKE_BUILD_TYPE=RelWithDebInfo \
+    -DMLIR_DIR=$llvm_dir/lib/cmake/mlir \
+    -DLLVM_EXTERNAL_LIT=$llvm_dir/bin/llvm-lit
 cmake --build build --parallel $(nproc)
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -1,7 +1,13 @@
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
-set(gc_opt_libs ${dialect_libs} ${conversion_libs} MLIROptLib)
+
+set(gc_opt_libs 
+        ${dialect_libs}
+        ${conversion_libs} 
+        MLIROptLib 
+        GCPasses)
 add_llvm_executable(gc-opt gc-opt.cpp)
+
 target_link_libraries(gc-opt PRIVATE ${gc_opt_libs})
 llvm_update_compile_flags(gc-opt)
 mlir_check_all_link_libraries(gc-opt)
diff --git a/src/gc-opt.cpp b/src/gc-opt.cpp
@@ -17,12 +17,15 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+#include "gc-dialects/Passes.h"
 #include "mlir/InitAllDialects.h"
 #include "mlir/InitAllPasses.h"
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
 
 int main(int argc, char *argv[]) {
   mlir::registerAllPasses();
+  mlir::gc::registerGraphCompilerPasses();
+
   mlir::DialectRegistry registry;
   mlir::registerAllDialects(registry);
   return mlir::asMainReturnCode(mlir::MlirOptMain(
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
@@ -10,3 +10,28 @@ add_library(graph_compiler_static STATIC ${GC_LIB_SOURCES})
 target_include_directories(graph_compiler_static PUBLIC ${GC_LIB_INCLUDES})
 
 add_subdirectory(dnnl)
+
+configure_lit_site_cfg(
+        ${CMAKE_CURRENT_SOURCE_DIR}/lit.site.cfg.py.in
+        ${CMAKE_CURRENT_BINARY_DIR}/lit.site.cfg.py
+        MAIN_CONFIG
+        ${CMAKE_CURRENT_SOURCE_DIR}/lit.cfg.py
+)
+
+set(GC_OPT_TEST_DEPENDS
+        FileCheck count not
+        # mlir-gen
+        gc-opt
+        )
+
+add_lit_testsuite(gc-check "Running the regression tests"
+        ${CMAKE_CURRENT_BINARY_DIR}
+        DEPENDS ${GC_OPT_TEST_DEPENDS}
+        )
+
+# Looks that this property is suitable for IDE
+# TODO: Check is this fine for IDE
+set_target_properties(gc-check PROPERTIES FOLDER "Tests")
+
+add_lit_testsuites(GC_OPT ${CMAKE_CURRENT_SOURCE_DIR} DEPENDS ${GC_OPT_TEST_DEPENDS})
+
diff --git a/test/gc-dialects/Linlagx/linalg-generalize.mlir b/test/gc-dialects/Linlagx/linalg-generalize.mlir
@@ -0,0 +1,73 @@
+// RUN: gc-opt %s --split-input-file --linalg-generalize-named-ops | FileCheck %s
+
+
+// CHECK: #[[$MAP0:.+]] = affine_map<(d0, d1) -> (d1, d0)>
+// CHECK: #[[$MAP1:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK: #[[$MAP2:.+]] = affine_map<(d0, d1) -> ()>
+// CHECK: #[[$MAP3:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK: #[[$MAP4:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
+// CHECK: #[[$MAP5:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK: #[[$MAP6:.+]] = affine_map<(d0, d1) -> (d1)>
+// CHECK: #[[$MAP7:.+]] = affine_map<(d0, d1) -> (0, d1)>
+#map = affine_map<(d0, d1) -> (d1)>
+#map1 = affine_map<(d0, d1) -> (0, d1)>
+#map2 = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: mlp
+// CHECK-SAME: %[[ARG0:.+]]: tensor<8192x16384xf32>, %[[ARG1:.+]]: tensor<8192xf32>, %[[ARG2:.+]]: tensor<1x128x128xf32>
+func.func @mlp(%arg0: tensor<8192x16384xf32>, %arg1: tensor<8192xf32>, %arg2: tensor<1x128x128xf32>) -> tensor<1x8192xf32> {
+  // CHECK: %[[CST:.+]] = arith.constant 0.000000e+00 : f32
+  %cst = arith.constant 0.000000e+00 : f32
+  
+  // CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %arg2
+  // CHECK-SAME{literal}: [[0], [1, 2]] : tensor<1x128x128xf32> into tensor<1x16384xf32>
+  %collapsed = tensor.collapse_shape %arg2 [[0], [1, 2]] : tensor<1x128x128xf32> into tensor<1x16384xf32>
+  
+  // CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<16384x8192xf32>
+  %0 = tensor.empty() : tensor<16384x8192xf32>
+  
+  // CHECK: %[[TRANSPOSED:.+]] = linalg.generic
+  // CHECK-SAME: {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel"]} 
+  // CHECK-SAME: ins(%arg0 : tensor<8192x16384xf32>) outs(%[[EMPTY]] : tensor<16384x8192xf32>) {
+  // CHECK: ^bb0(%[[IN:.+]]: f32, %[[OUT:.+]]: f32):
+  // CHECK:   linalg.yield %[[IN]] : f32
+  // CHECK: } -> tensor<16384x8192xf32>
+  %transposed = linalg.transpose ins(%arg0 : tensor<8192x16384xf32>) outs(%0 : tensor<16384x8192xf32>) permutation = [1, 0] 
+  
+  // CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x8192xf32>
+  %1 = tensor.empty() : tensor<1x8192xf32>
+  
+  // CHECK: %[[FILLED:.+]] = linalg.generic
+  // CHECK-SAME: {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel"]}
+  // CHECK-SAME: ins(%[[CST]] : f32) outs(%[[EMPTY]] : tensor<1x8192xf32>)
+  // CHECK: ^bb0(%[[IN:.+]]: f32, %[[OUT:.+]]: f32):
+  // CHECK:   linalg.yield %[[IN]] : f32
+  // CHECK: } -> tensor<1x8192xf32>
+  %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<1x8192xf32>) -> tensor<1x8192xf32>
+
+  // CHECK: %[[MMRESULT:.+]] = linalg.generic
+  // CHECK-SAME: {indexing_maps = [#map3, #map4, #map5], iterator_types = ["parallel", "parallel", "reduction"]}
+  // CHECK-SAME: ins(%[[COLLAPSE]], %[[TRANSPOSED]] : tensor<1x16384xf32>, tensor<16384x8192xf32>) outs(%[[FILLED]] : tensor<1x8192xf32>)
+  // CHECK: ^bb0(%[[IN0:.+]]: f32, %[[IN1:.+]]: f32, %[[OUT:.+]]: f32):
+  // CHECK:   %[[MUL:.+]] = arith.mulf %[[IN0]], %[[IN1]] : f32
+  // CHECK:   %[[ADD:.+]] = arith.addf %[[OUT]], %[[MUL]] : f32
+  // CHECK:   linalg.yield %[[ADD]] : f32
+  // CHECK: } -> tensor<1x8192xf32>
+  %3 = linalg.matmul ins(%collapsed, %transposed : tensor<1x16384xf32>, tensor<16384x8192xf32>) outs(%2 : tensor<1x8192xf32>) -> tensor<1x8192xf32>
+  
+  // CHECK: %[[BIAS:.+]] = linalg.generic
+  // CHECK-SAME: {indexing_maps = [#map6, #map7, #map1], iterator_types = ["parallel", "parallel"]}
+  // CHECK-SAME: ins(%arg1, %[[MMRESULT]] : tensor<8192xf32>, tensor<1x8192xf32>) outs(%[[EMPTY]] : tensor<1x8192xf32>) {
+  // CHECK: ^bb0(%[[IN0:.+]]: f32, %[[IN1:.+]]: f32, %[[OUT:.+]]: f32):
+  // CHECK:   %[[ADD:.+]] = arith.addf %[[IN0]], %[[IN1]] : f32
+  // CHECK:   linalg.yield %[[ADD]] : f32
+  // CHECK: } -> tensor<1x8192xf32>
+  %4 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel"]} ins(%arg1, %3 : tensor<8192xf32>, tensor<1x8192xf32>) outs(%1 : tensor<1x8192xf32>) {
+  ^bb0(%in: f32, %in_0: f32, %out: f32):
+    %5 = arith.addf %in, %in_0 : f32
+    linalg.yield %5 : f32
+  } -> tensor<1x8192xf32>
+
+  // CHECK: return %[[BIAS]] : tensor<1x8192xf32>
+  return %4 : tensor<1x8192xf32>
+}
+// -----
diff --git a/test/lit.cfg.py b/test/lit.cfg.py
@@ -0,0 +1,52 @@
+# -*- Python -*-
+
+import os
+import lit.formats
+import lit.util
+
+from lit.llvm import llvm_config
+
+# from lit.llvm.subst import ToolSubst
+# from lit.llvm.subst import FindTool
+
+# Configuration file for the 'lit' test runner.
+
+# name: The name of this test suite.
+config.name = "GRAPH_COMPILER_OPT"
+
+config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell)
+
+# suffixes: A list of file extensions to treat as test files.
+config.suffixes = [".mlir"]
+
+# test_source_root: The root path where tests are located.
+config.test_source_root = os.path.dirname(__file__)
+
+# test_exec_root: The root path where tests should be run.
+config.test_exec_root = os.path.join(config.gc_obj_root, "test")
+
+config.substitutions.append(("%PATH%", config.environment["PATH"]))
+config.substitutions.append(("%shlibext", config.llvm_shlib_ext))
+config.substitutions.append(("%llvmlibdir", config.llvm_lib_dir))
+config.substitutions.append(("%gclibdir", config.gc_obj_root + "/lib/"))
+
+llvm_config.with_system_environment(["HOME", "INCLUDE", "LIB", "TMP", "TEMP"])
+
+llvm_config.use_default_substitutions()
+
+# excludes: A list of directories to exclude from the testsuite. The 'Inputs'
+# subdirectories contain auxiliary inputs for various tests in their parent
+# directories.
+config.excludes = []
+
+# test_exec_root: The root path where tests should be run.
+config.test_exec_root = os.path.join(config.gc_obj_root, "test")
+config.gc_tools_dir = os.path.join(config.gc_obj_root, "src")
+
+# Tweak the PATH to include the tools dir.
+llvm_config.with_environment("PATH", config.llvm_tools_dir, append_path=True)
+
+tool_dirs = [config.gc_tools_dir, config.llvm_tools_dir]
+tools = ["gc-opt"]
+
+llvm_config.add_tool_substitutions(tools, tool_dirs)
diff --git a/test/lit.site.cfg.py.in b/test/lit.site.cfg.py.in
@@ -0,0 +1,38 @@
+@LIT_SITE_CFG_IN_HEADER@
+
+import sys
+
+config.llvm_src_root = "@LLVM_SOURCE_DIR@"
+config.llvm_obj_root = "@LLVM_BINARY_DIR@"
+config.llvm_tools_dir = "@LLVM_TOOLS_DIR@"
+config.llvm_lib_dir = "@LLVM_LIBS_DIR@"
+config.llvm_shlib_dir = "@SHLIBDIR@"
+config.llvm_shlib_ext = "@SHLIBEXT@"
+config.llvm_exe_ext = "@EXEEXT@"
+
+config.lit_tools_dir = "@LLVM_LIT_TOOLS_DIR@"
+config.python_executable = "@PYTHON_EXECUTABLE@"
+config.gold_executable = "@GOLD_EXECUTABLE@"
+config.ld64_executable = "@LD64_EXECUTABLE@"
+config.enable_shared = @ENABLE_SHARED@
+config.enable_assertions = @ENABLE_ASSERTIONS@
+config.targets_to_build = "@TARGETS_TO_BUILD@"
+config.native_target = "@LLVM_NATIVE_ARCH@"
+config.llvm_bindings = "@LLVM_BINDINGS@".split(' ')
+config.host_os = "@HOST_OS@"
+config.host_cc = "@HOST_CC@"
+config.host_cxx = "@HOST_CXX@"
+# Note: ldflags can contain double-quoted paths, so must use single quotes here.
+config.host_ldflags = '@HOST_LDFLAGS@'
+config.llvm_use_sanitizer = "@LLVM_USE_SANITIZER@"
+config.llvm_host_triple = '@LLVM_HOST_TRIPLE@'
+config.host_arch = "@HOST_ARCH@"
+config.gc_src_root = "@CMAKE_SOURCE_DIR@"
+config.gc_obj_root = "@CMAKE_BINARY_DIR@"
+
+
+import lit.llvm
+lit.llvm.initialize(lit_config, config)
+
+# Let the main config do the real work.
+lit_config.load_config(config, "@CMAKE_SOURCE_DIR@/test/lit.cfg.py")
