#!/bin/bash -e
#
# Apache v2 license
# Copyright (C) 2023 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
#

OPTION=${1:-inference_throughput_amx_bfloat16}
PLATFORM=${PLATFORM:-SPR}
CORES_PER_INSTANCE=${CORES_PER_INSTANCE:-}
INSTANCE_NUMBER=${INSTANCE_NUMBER:-}
TOPOLOGY="bert_large"
PRECISION=${PRECISION:-avx_fp32}
DATA_TYPE=${DATA_TYPE:-real}
ONEDNN_VERBOSE=${ONEDNN_VERBOSE:-0}
FUNCTION=$(echo ${OPTION}|cut -d_ -f1)
MODE=$(echo ${OPTION}|cut -d_ -f2)
PRECISION=$(echo ${OPTION}|cut -d_ -f3-4)
CASE_TYPE=$(echo ${OPTION}|cut -d_ -f5)
CUSTOMER_ENV=${CUSTOMER_ENV}

if [ -n "$CASE_TYPE" ] && [ "$CASE_TYPE" == "pkm" ]; then
    if [ "$MODE" == "throughput" ]; then
        EVENT_TRACE_PARAMS="roi,Running evaluation,Evaluating: 100%"
    fi
fi

if [ "$FUNCTION" == "inference" ]; then
    BATCH_SIZE=${BATCH_SIZE:-1}
    MAX_SEQ_LENGTH=${MAX_SEQ_LENGTH:-384}
    WARMUP_STEPS=${WARMUP_STEPS:-15}
    STEPS=${STEPS:-30}
    WEIGHT_SHARING=${WEIGHT_SHARING:-False}
else
    BATCH_SIZE=${BATCH_SIZE:-28}
    MAX_SEQ_LENGTH=512
    WARMUP_STEPS=${WARMUP_STEPS:-0}
    STEPS=51
    NNODES=${NNODES:-1}
    DISTRIBUTED=${DISTRIBUTED:-False}
    CCL_WORKER_COUNT=${CCL_WORKER_COUNT:-4}
fi

# Logs Setting
DIR="$( cd "$( dirname "$0" )" &> /dev/null && pwd )"
. "$DIR/../../script/overwrite.sh"

# Workload Setting
if [ "${FUNCTION}" == "inference" ]; then
    WORKLOAD_PARAMS=(MODE PLATFORM TOPOLOGY PRECISION FUNCTION DATA_TYPE WARMUP_STEPS STEPS BATCH_SIZE CORES_PER_INSTANCE INSTANCE_NUMBER WEIGHT_SHARING ONEDNN_VERBOSE ENABLE_PROFILING MAX_SEQ_LENGTH CUSTOMER_ENV)
elif [ "${DISTRIBUTED}" == "True" ]; then
    WORKLOAD_PARAMS=(MODE PLATFORM TOPOLOGY PRECISION FUNCTION DATA_TYPE WARMUP_STEPS STEPS BATCH_SIZE NNODES CORES_PER_INSTANCE INSTANCE_NUMBER WEIGHT_SHARING DISTRIBUTED CCL_WORKER_COUNT ONEDNN_VERBOSE ENABLE_PROFILING MAX_SEQ_LENGTH CUSTOMER_ENV)
else
    WORKLOAD_PARAMS=(MODE PLATFORM TOPOLOGY PRECISION FUNCTION DATA_TYPE WARMUP_STEPS STEPS BATCH_SIZE CORES_PER_INSTANCE INSTANCE_NUMBER WEIGHT_SHARING DISTRIBUTED ONEDNN_VERBOSE ENABLE_PROFILING MAX_SEQ_LENGTH CUSTOMER_ENV)
fi

# Docker Setting
if [[ "$WORKLOAD" = *pdt ]]; then
    DOCKER_IMAGE="$DIR/pdt/Dockerfile.1.${FUNCTION}"
elif [[ -e $DIR/build_ext.sh ]]; then
    DOCKER_IMAGE="$DIR/Dockerfile.1.${FUNCTION}_24.04"
else
    DOCKER_IMAGE="$DIR/Dockerfile.1.${FUNCTION}"
fi

DOCKER_OPTIONS="--privileged -e MODE=${MODE} -e PLATFORM=${PLATFORM} \
                -e TOPOLOGY=${TOPOLOGY} -e PRECISION=${PRECISION} \
                -e FUNCTION=${FUNCTION} -e DATA_TYPE=${DATA_TYPE} \
                -e STEPS=${STEPS} -e BATCH_SIZE=${BATCH_SIZE} -e WARMUP_STEPS=${WARMUP_STEPS} \
                -e CORES_PER_INSTANCE=${CORES_PER_INSTANCE} \
                -e INSTANCE_NUMBER=${INSTANCE_NUMBER} -e ONEDNN_VERBOSE=${ONEDNN_VERBOSE} \
                -e ENABLE_PROFILING=${ENABLE_PROFILING} -e MAX_SEQ_LENGTH=${MAX_SEQ_LENGTH} \
                -e WEIGHT_SHARING=${WEIGHT_SHARING} \
                -e CUSTOMER_ENV=${CUSTOMER_ENV}"


# Kubernetes Setting
RECONFIG_OPTIONS="-DK_TOPOLOGY=${TOPOLOGY} -DK_MODE=${MODE} -DK_PLATFORM=${PLATFORM} \
                  -DK_PRECISION=${PRECISION} -DK_FUNCTION=${FUNCTION} \
                  -DK_DATA_TYPE=${DATA_TYPE} -DK_STEPS=${STEPS} \
                  -DK_BATCH_SIZE=${BATCH_SIZE} -DK_WARMUP_STEPS=${WARMUP_STEPS} \
                  -DK_CORES_PER_INSTANCE=${CORES_PER_INSTANCE} \
                  -DK_INSTANCE_NUMBER=${INSTANCE_NUMBER} \
                  -DK_ONEDNN_VERBOSE=${ONEDNN_VERBOSE} \
                  -DK_ENABLE_PROFILING=${ENABLE_PROFILING} \
                  -DK_MAX_SEQ_LENGTH=${MAX_SEQ_LENGTH} \
                  -DK_WEIGHT_SHARING=${WEIGHT_SHARING} \
                  -DK_CUSTOMER_ENV=${CUSTOMER_ENV} \
                  -DK_DOCKER_IMAGE=${DOCKER_IMAGE}"

if [ "${FUNCTION}" == "training" ]; then
    DOCKER_OPTIONS+=" -e DISTRIBUTED=${DISTRIBUTED}"
    RECONFIG_OPTIONS+=" -DK_DISTRIBUTED=${DISTRIBUTED}"
    if [ "${DISTRIBUTED}" == "True" ]; then
        DOCKER_OPTIONS+=" -e NNODES=${NNODES} -e CCL_WORKER_COUNT=${CCL_WORKER_COUNT}"
        RECONFIG_OPTIONS+=" -DK_NNODES=${NNODES} -DK_CCL_WORKER_COUNT=${CCL_WORKER_COUNT}"
    fi
fi
                  
JOB_FILTER="job-name=bertlarge-pytorch-xeon-public-benchmark"

. "$DIR/../../script/validate.sh"
