netron/tools/mlir

#!/bin/bash

set -e
pushd $(cd $(dirname ${0})/..; pwd) > /dev/null

src_dir=./third_party/source/mlir

entries=(
    "https://github.com/llvm/llvm-project.git|main|${src_dir}/llvm-project|mlir"
    "https://github.com/openxla/stablehlo.git|main|${src_dir}/stablehlo|stablehlo"
    "https://github.com/openxla/shardy.git|main|${src_dir}/shardy|shardy"
    "https://github.com/openxla/xla.git|main|${src_dir}/xla|xla/mlir_hlo"
    "https://github.com/onnx/onnx-mlir.git|main|${src_dir}/onnx-mlir|"
    "https://github.com/llvm/torch-mlir.git|main|${src_dir}/torch-mlir|"
    "https://github.com/triton-lang/triton|main|${src_dir}/triton|"
    "https://github.com/tensorflow/tensorflow.git|master|${src_dir}/tensorflow|tensorflow/compiler/mlir"
    "https://github.com/tensorflow/runtime.git|master|${src_dir}/runtime|"
    "https://github.com/NVIDIA/TensorRT-Incubator.git|main|${src_dir}/TensorRT-Incubator|mlir-tensorrt"
    "https://github.com/iree-org/iree.git|main|${src_dir}/iree|"
    "https://github.com/monellz/FlashTensor.git|main|${src_dir}/FlashTensor|"
    "https://github.com/plaidml/plaidml.git|plaidml-v1|${src_dir}/plaidml|pmlc"
    "https://github.com/sophgo/tpu-mlir.git|master|${src_dir}/tpu-mlir|include"
    "https://github.com/spcl/mlir-dace.git|main|${src_dir}/mlir-dace|"
    "https://github.com/woxjro/lltz.git|master|${src_dir}/lltz|"
    "https://github.com/pengmai/lagrad.git|main|${src_dir}/lagrad|include"
    "https://github.com/llvm/clangir.git|main|${src_dir}/clangir|clang/include/clang/CIR"
    "https://github.com/ROCm/rocMLIR.git|develop|${src_dir}/rocMLIR|mlir/include/mlir/Dialect"
)

clean() {
    echo "mlir clean"
    for entry in "${entries[@]}"; do
        IFS='|' read -r url branch dir root <<< "${entry}"
        rm -rf "${dir}"
    done
}

sync() {
    echo "mlir sync"
    for entry in "${entries[@]}"; do
        IFS='|' read -r url branch dir root <<< "${entry}"
        if [ ! -d "${dir}" ]; then
            if [ -n "${root}" ]; then
                git clone --quiet --depth=1 --branch "${branch}" --single-branch --filter=blob:none --sparse "${url}" "${dir}" >/dev/null 2>&1
                git -C "${dir}" sparse-checkout set "${root}" >/dev/null 2>&1
            else
                git clone --quiet --depth=1 --branch "${branch}" --single-branch "${url}" "${dir}"
            fi
        else
            if [ -n "${root}" ]; then
                git -C "${dir}" sparse-checkout set "${root}" >/dev/null 2>&1
            fi
            git -C "${dir}" fetch --quiet --depth=1 origin "${branch}"
            git -C "${dir}" reset --quiet --hard FETCH_HEAD >/dev/null 2>&1
            git -C "${dir}" gc --quiet --prune=all
        fi
    done
    mkdir -p "${src_dir}/_/llvm-project/mlir/include/mlir/Interfaces"
    curl --silent --show-error --location --output "${src_dir}/_/llvm-project/mlir/include/mlir/Interfaces/CopyOpInterface.td" "https://raw.githubusercontent.com/llvm/llvm-project/llvmorg-14.0.0/mlir/include/mlir/Interfaces/CopyOpInterface.td"
    mkdir -p "${src_dir}/_/llvm-project/mlir/include/mlir/Dialect/Arithmetic/IR"
    curl --silent --show-error --location --output "${src_dir}/_/llvm-project/mlir/include/mlir/Dialect/Arithmetic/IR/ArithmeticBase.td" "https://raw.githubusercontent.com/llvm/llvm-project/llvmorg-14.0.0/mlir/include/mlir/Dialect/Arithmetic/IR/ArithmeticBase.td"
    mkdir -p "${src_dir}/_/llvm-project/mlir/include/mlir/Dialect/OpenACC"
    echo '#ifndef ACCCOMMON_TD
#define ACCCOMMON_TD
#endif // ACCCOMMON_TD' > "${src_dir}/_/llvm-project/mlir/include/mlir/Dialect/OpenACC/AccCommon.td"
    mkdir -p "${src_dir}/_/llvm-project/mlir/include/mlir/Dialect/OpenMP"
    echo '#ifndef OMPCOMMON_TD
#define OMPCOMMON_TD
#endif // OMPCOMMON_TD' > "${src_dir}/_/llvm-project/mlir/include/mlir/Dialect/OpenMP/OmpCommon.td"
    mkdir -p "${src_dir}/_/llvm-project/mlir/include/mlir/Dialect/Linalg/IR"
    echo '#ifndef LINALG_NAMED_STRUCTURED_OPS_YAMLGEN_TD
#define LINALG_NAMED_STRUCTURED_OPS_YAMLGEN_TD
#endif // LINALG_NAMED_STRUCTURED_OPS_YAMLGEN_TD' > "${src_dir}/_/llvm-project/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yamlgen.td"
}

schema() {
    echo "mlir schema"
    node ./tools/mlir-script.js schema
}

test() {
    echo "mlir test"
    node ./tools/mlir-script.js test "$@"
}

while [ "$#" != 0 ]; do
    command="$1" && shift
    case "${command}" in
        "clean") clean;;
        "sync") sync;;
        "schema") schema;;
        "test") test "$@";;
    esac
done