#!/bin/bash

# 设置工作目录
WORK_DIR=$(pwd)
PROJECT_NAME="torch-mlir"
PROJECT_REPO="https://github.com/llvm/torch-mlir.git"
COMMIT_ID="7e7af670802d99cacdaf26e6e37249d544e4896e"

# 创建并进入临时目录
mkdir -p "$WORK_DIR"/"$PROJECT_NAME"
cd "$WORK_DIR"/"$PROJECT_NAME"

# 克隆torch-mlir仓库
echo "Cloning torch-mlir repository..."
git clone "$PROJECT_REPO" --depth=1
cd "$PROJECT_NAME"
git remote set-branches origin "$COMMIT_ID"
git fetch --depth 1 origin "$COMMIT_ID"
git checkout "$COMMIT_ID"
cd ..

# 进入项目目录
cd "$PROJECT_NAME"

# 初始化子模块
echo "Initializing submodules..."
git submodule update --init --progress

# 配置编译选项
echo "Configuring build..."
cmake -GNinja -Bbuild \
  -DCMAKE_BUILD_TYPE=Release \
  -DPython3_FIND_VIRTUALENV=ONLY \
  -DLLVM_ENABLE_PROJECTS=mlir \
  -DLLVM_EXTERNAL_PROJECTS="torch-mlir" \
  -DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR="$PWD" \
  -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
  -DLLVM_TARGETS_TO_BUILD=host \
  -DTORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS=ON \
  -DTORCH_MLIR_ENABLE_JIT_IR_IMPORTER=ON \
  externals/llvm-project/llvm

# 编译torch-mlir
echo "Building torch-mlir..."
cmake --build build --target tools/torch-mlir/all -j32

# install torch-mlir python package
echo "Installing torch-mlir python package..."
TORCH_MLIR_CMAKE_BUILD_DIR=build/ TORCH_MLIR_CMAKE_ALREADY_BUILT=1 python setup.py bdist_wheel
pip install dist/torch_mlir-0.0.1-*.whl

echo "Build completed."

# 返回原工作目录
cd "$WORK_DIR"