#!/bin/bash

# SPDX-FileCopyrightText: Copyright (c) 2019-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.

NRDEVICES=$(grep -c ^processor /proc/cpuinfo | sed 's/^0$/1/')
if modinfo zram | grep -q ' zram_num_devices:' 2>/dev/null; then
	MODPROBE_ARGS="zram_num_devices=${NRDEVICES}"
elif modinfo zram | grep -q ' num_devices:' 2>/dev/null; then
	MODPROBE_ARGS="num_devices=${NRDEVICES}"
else
	exit 1
fi
modprobe zram "${MODPROBE_ARGS}"

# Calculate memory to use for zram (1/2 of ram)
totalmem=$(LC_ALL=C free | grep -e "^Mem:" | sed -e 's/^Mem: *//' -e 's/  *.*//')
mem=$((("${totalmem}" / 2 / "${NRDEVICES}") * 1024))

# initialize the devices
for i in $(seq "${NRDEVICES}"); do
	DEVNUMBER=$((i - 1))
	echo "${mem}" > /sys/block/zram"${DEVNUMBER}"/disksize
	mkswap /dev/zram"${DEVNUMBER}"
	swapon -p 5 /dev/zram"${DEVNUMBER}"
done