Integrating alp with iree-llvm-sandbox

CMakeLists.txt

@@ -32,6 +32,13 @@ set(MLIR_MAIN_SRC_DIR ${LLVM_MAIN_SRC_DIR}/../mlir)
 set(MLIR_INCLUDE_DIR ${LLVM_MAIN_SRC_DIR}/../mlir/include)
 set(MLIR_TABLEGEN_OUTPUT_DIR ${CMAKE_BINARY_DIR}/tools/mlir/include)
 
+if (SANDBOX_ENABLE_ALP)
+  # ALP includes
+  add_compile_definitions("SANDBOX_ENABLE_ALP")
+  include_directories(experimental/alp/include)
+  include_directories(${CMAKE_BINARY_DIR}/tools/sandbox/experimental/alp/include)
+endif()
+
 list(APPEND CMAKE_MODULE_PATH ${MLIR_MAIN_SRC_DIR}/cmake/modules)
 list(APPEND CMAKE_MODULE_PATH ${LLVM_MAIN_SRC_DIR}/cmake)
 set(MLIR_TABLEGEN_EXE mlir-tblgen)
@@ -80,3 +87,8 @@ add_subdirectory(lib)
 add_subdirectory(python)
 add_subdirectory(test)
 add_subdirectory(tools)
+if (SANDBOX_ENABLE_ALP)
+# ALP cmake files
+add_subdirectory(experimental)
+endif()
+

configure.py

@@ -23,6 +23,11 @@ def parse_arguments():
                       help="Build with ENABLE_LLD=ON (optional)",
                       dest="enable_lld",
                       default = False)
+  parser.add_argument("--enable-alp", 
+                      help="Build with SANDBOX_ENABLE_ALP=ON (optional)",
+                      dest="enable_alp",
+                      action="store_false",
+                      default = False)
   parser.add_argument("--no-ccache",
                       help="Disables ccache (if available)",
                       dest="enable_ccache",
@@ -102,6 +107,10 @@ def main(args):
     else:
       print("WARNING: LLD (ld.lld) not found on path. Configure may fail.")
 
+  # Optionally enable Alp
+  if args.enable_alp:
+    llvm_configure_args.append("-DSANBOX_ENABLE_ALP")
+
   # Detect ccache.
   if args.enable_ccache:
     ccache_path = shutil.which("ccache")

experimental/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(alp)

experimental/alp/CMakeLists.txt

@@ -0,0 +1,14 @@
+get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
+
+set(LLVM_LINK_COMPONENTS
+  Core
+  Support
+  nativecodegen
+  native
+  OrcJIT
+  )
+
+include_directories(include/)
+add_subdirectory(include)
+add_subdirectory(lib)

experimental/alp/alp/compile_op.py

@@ -0,0 +1,219 @@
+#  Part of the LLVM Project, 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
+import os
+import tempfile
+import shutil
+from pathlib import Path
+from .utils import print_command, run_and_save, run_command, add_extension
+from .library.blas import gemm
+
+def build_main_obj(benchmark_prog, m, n, k, op, reps, mktmp_fn):
+    benchmark_prog = benchmark_prog.replace("_M_", str(m))
+    benchmark_prog = benchmark_prog.replace("_K_", str(k))
+    benchmark_prog = benchmark_prog.replace("_N_", str(n))
+    benchmark_prog = benchmark_prog.replace("__OP__", op)
+    benchmark_prog = benchmark_prog.replace("_REPS_", str(reps))
+
+    main_mlir = mktmp_fn("test.mlir")
+    main_mlir_lowered = mktmp_fn("test.llvm.mlir")
+    main_llvm = mktmp_fn("test.ll")
+    main_obj = mktmp_fn("test.o")
+
+    f = open(main_mlir, "w")
+    f.write(benchmark_prog)
+    f.close()
+
+    # main program
+    cmd = ["mlir-opt"]
+    cmd.append(main_mlir)
+    cmd.append("--linalg-bufferize")
+    cmd.append("--std-bufferize")
+    cmd.append("--tensor-constant-bufferize")
+    cmd.append("--tensor-bufferize")
+    cmd.append("--func-bufferize")
+    cmd.append("-convert-linalg-to-affine-loops")
+    cmd.append("-lower-affine")
+    cmd.append("-convert-scf-to-std")
+    cmd.append("-convert-memref-to-llvm")
+    cmd.append("-convert-std-to-llvm")
+    cmd.append("-reconcile-unrealized-casts")
+    cmd.append(f"> {main_mlir_lowered}")
+    run_command(cmd)
+    print_command(cmd)
+
+    cmd = ["mlir-translate"]
+    cmd.append("--mlir-to-llvmir")
+    cmd.append(f"{main_mlir_lowered}")
+    cmd.append(f"> {main_llvm}")
+    run_command(cmd)
+
+    cmd = ["llc"]
+    cmd.append(f"{main_llvm}")
+    cmd.append("-O3")
+    cmd.append("-filetype=obj")
+    cmd.append(f"-o {main_obj}")
+    run_command(cmd)
+
+
+def apply(transform_list, op_mlir_file, verbosity_level):
+    cmd = ["$IREE_LLVM_SANDBOX_BUILD_DIR/bin/mlir-proto-opt "]
+    for t in transform_list:
+      if not t:
+        continue
+      if type(t) is tuple:
+        (l, ext) = t
+        if l >= verbosity_level:
+          run_and_save(cmd, op_mlir_file, add_extension(op_mlir_file, ext))
+      else:
+        cmd.append(t)
+
+    output = add_extension(op_mlir_file, "llvm")
+    run_and_save(cmd, op_mlir_file, output)
+    return output
+
+def SaveIR(x, ext):
+  return (x, ext)
+
+def build_operator_obj(op_prog, m, n, k, op, option_list, mktmp_fn, verbosity_level=0):
+    op_prog = op_prog.replace("_M_", str(m))
+    op_prog = op_prog.replace("_K_", str(k))
+    op_prog = op_prog.replace("_N_", str(n))
+
+    op_mlir = mktmp_fn(f"{op}.mlir")
+    f = open(f"{op_mlir}", "w")
+    f.write(op_prog)
+    f.close()
+
+    # Transformation options
+    tile_sizes = option_list["tile_sizes"]
+    reorder_tile_sizes = option_list["reorder_tile_sizes"]
+    register_tile_sizes = option_list["register_tile_sizes"]
+    reorder_register_tile_sizes = option_list["reorder_register_tile_sizes"]
+    hoist_packing = option_list['hoist_packing']
+    split_vector_transfer = option_list['split_vector_transfers_to']
+    extract_micro_kernel = option_list['extract_micro_kernel']
+    modulo_scheduling = option_list['modulo_scheduling']
+
+    Canonicalize = " --canonicalize --cse"
+    CodegenDriver = "--linalg-tensor-codegen-driver=\"anchor-func=gemm anchor-op=linalg.generic" 
+
+    # Transformations
+    OuterTiling = CodegenDriver + f" tile-sizes={tile_sizes} tile-interchange={reorder_tile_sizes}\"" + Canonicalize
+
+    InnerTiling = CodegenDriver + f" tile-sizes={register_tile_sizes} tile-interchange={reorder_register_tile_sizes}" + \
+                                  f" pad pack-paddings=1,1,0 hoist-paddings={hoist_packing} \"" + Canonicalize
+
+    DecomposeToLowerDimensionalNamedOp = CodegenDriver + " decompose-to-lower-dim\"" + Canonicalize
+
+    Vectorize = CodegenDriver + " vectorize vectorize-padding\"" + Canonicalize
+
+    Bufferize = "--linalg-bufferization-driver" + Canonicalize
+
+    LowerVector = "--linalg-vector-lowering=\"max-transfer-rank=1 " +\
+                 f" split-transfers={split_vector_transfer}" +\
+                  " lower-vector-transpose-to=eltwise" +\
+                  " lower-vector-multi-reduction-to=innerparallel" +\
+                  " lower-vector-contraction-to=outerproduct" +\
+                  " unroll-vector-transfers=true" 
+
+    LowerVectorStage = lambda stage : LowerVector+f" lower-vector-stage={stage}\"" + Canonicalize
+
+    ExtractKernel = "--alp-extract-kernel" + Canonicalize if extract_micro_kernel else ""
+    ModuloScheduling = "--alp-modulo-scheduling"  if modulo_scheduling else "" # TODO: Order is not preserved if I canonicalize
+
+    LowerToLLVM = "--convert-vector-to-scf " +\
+                  "--convert-linalg-to-loops " +\
+                  "--canonicalize " +\
+                  "--lower-affine " +\
+                  "--convert-scf-to-std " +\
+                  "--convert-linalg-to-llvm " +\
+                  "--convert-vector-to-llvm " +\
+                  "--convert-math-to-llvm " +\
+                  "--convert-memref-to-llvm " +\
+                  "--convert-std-to-llvm " +\
+                  "--canonicalize " +\
+                  "--cse " +\
+                  "--reconcile-unrealized-casts " 
+
+
+    TransformList = [OuterTiling, 
+                     InnerTiling, 
+                     SaveIR(4, "tile"),
+                     DecomposeToLowerDimensionalNamedOp, 
+                     Vectorize,
+                     SaveIR(4, "vectorize"),
+                     Bufferize,
+                     SaveIR(4, "bufferize"),
+                     LowerVectorStage(0),
+                     SaveIR(4, "lower_vector"),
+                     ExtractKernel,
+                     ModuloScheduling,
+                     SaveIR(4, "micro_kernel"),
+                     LowerVectorStage(1),
+                     LowerVectorStage(2),
+                     LowerVectorStage(3),
+                     LowerVectorStage(4),
+                     LowerVectorStage(5),
+                     LowerVectorStage(6),
+                     SaveIR(4, "micro_kernel_final"),
+                     LowerToLLVM]
+
+    op_llvm_mlir = apply(TransformList, op_mlir, verbosity_level)
+
+    out = run_command(["$IREE_LLVM_SANDBOX_BUILD_DIR/bin/mlir-translate --mlir-to-llvmir " + op_llvm_mlir])
+    op_llvm = mktmp_fn(f"{op}.ll")
+    f = open(f"{op_llvm}", "w")
+    f.write(out)
+    f.close()
+
+    op_obj = mktmp_fn(f"{op}.o")
+    op_asm = mktmp_fn(f"{op}.s")
+
+    cmd = ["llc"]
+    cmd.append(op_llvm)
+    cmd.append("-O3")
+    cmd.append("-filetype=obj")
+    cmd.append(f"-o {op_obj}")
+    run_command(cmd)
+
+    cmd = ["llc"]
+    cmd.append(f"{op_llvm}")
+    cmd.append("-O3")
+    cmd.append("-filetype=asm")
+    cmd.append(f"-o {op_asm}")
+    run_command(cmd)
+
+def link_main(op, mktmp_fn):
+    out_bin = "exec_matmul"
+    main_obj = mktmp_fn("test.o")
+    op_obj = mktmp_fn(f"{op}.o")
+
+    cmd = ["clang++"]
+    cmd.append(f"{main_obj}")
+    cmd.append(f"{op_obj}")
+    cmd.append(f"-o {out_bin}")
+    cmd.append("-lmlir_c_runner_utils")
+    print_command(cmd)
+    run_command(cmd)
+
+def build_mlir(op, m, n, k, options):
+    verbose = ("verbosity_level" in options) and options["verbosity_level"] > 0
+    reps= 1
+    if options["reps"]:
+      reps = options["reps"]
+
+    if verbose:
+      Path("./tmp").mkdir(exist_ok=True)
+      tmp_dir_name = "./tmp" 
+      verbosity_level=options["verbosity_level"]
+    else:
+      tmp_dir = tempfile.TemporaryDirectory()
+      tmp_dir_name = tmp_dir.name
+      verbosity_level=0
+
+    (benchmark, op_mlir)= gemm(False)
+    mktmp = lambda x : os.path.join(tmp_dir_name, x)
+    build_main_obj(benchmark, m, n, k, op, reps, mktmp)
+    build_operator_obj(op_mlir, m, n, k, op, options, mktmp, verbosity_level)
+    link_main(op, mktmp)