/*
 * Copyright (c) 2026 Advanced Micro Devices, Inc.
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

/*
 * AMD/Intel RAPL energy counters exposed as an hwpmc(4) PMC class.
 *
 * Read-only, system-scope (PMC_MODE_SC), 64-bit counters reporting
 * microjoules.
 */

#include <sys/param.h>
#include <sys/bus.h>
#include <sys/callout.h>
#include <sys/malloc.h>
#include <sys/mutex.h>
#include <sys/pmc.h>
#include <sys/pmckern.h>
#include <sys/priv.h>
#include <sys/proc.h>
#include <sys/smp.h>
#include <sys/systm.h>

#include <machine/cpu.h>
#include <machine/cpufunc.h>
#include <machine/cputypes.h>
#include <machine/specialreg.h>

#include <x86/x86_var.h>

#include <dev/hwpmc/hwpmc_rapl.h>

/* Energy counters are per-package/core domains, not per CPU: DOMWIDE. */
#define	RAPL_CAPS	(PMC_CAP_READ | PMC_CAP_DOMWIDE)

/* Worst-case package watts for sizing the guard timer. */
#define	RAPL_GUARD_WATT		1000

/* Guard interval clamp band (ms). */
#define	RAPL_GUARD_MIN_MS	10
#define	RAPL_GUARD_MAX_MS	60000

struct rapl_event {
	enum pmc_event	re_ev;
	uint32_t	re_msr;
	uint32_t	re_unit;	/* unit shift: 1 tick = 1/2^unit J */
};

struct rapl_value {
	uint64_t	rv_prev;	/* last raw 32-bit MSR value	*/
	uint64_t	rv_accum;
	sbintime_t	rv_prev_time;
	bool		rv_primed;
};

struct rapl_cpu {
	struct pmc_hw	rc_hw[RAPL_MAX_NPMCS];
	struct rapl_value rc_value[RAPL_MAX_NPMCS];
	struct mtx	rc_mtx;
	int		rc_nalloc;	/* allocated RAPL PMCs on this CPU */
};

static struct rapl_cpu **rapl_pcpu;

static struct rapl_event rapl_events[RAPL_MAX_NPMCS];
static struct pmc_descr rapl_pmcdesc[RAPL_MAX_NPMCS];
static int rapl_npmcs;
static int rapl_ri;

static struct callout	rapl_guard_callout;
static sbintime_t	rapl_guard_sbt;
static int		rapl_nalloc;
static cpuset_t		rapl_cpus;	/* CPUs with an allocated RAPL PMC */

static struct mtx	rapl_alloc_mtx;

/* Convert energy ticks to microjoules without overflowing uint64_t. */
static uint64_t
rapl_raw_to_uj(uint64_t raw, uint32_t shift)
{
	uint64_t unit, whole, frac;

	unit = 1ULL << shift;
	whole = raw / unit;
	frac = raw % unit;
	return (whole * 1000000ULL + (frac * 1000000ULL) / unit);
}

/* Fold a 32-bit MSR reading into the 64-bit accumulator, can recover one wrap. */
static void
rapl_update_delta(struct rapl_value *val, uint64_t cur)
{
	sbintime_t now = sbinuptime();
	uint64_t diff;

	cur &= UINT32_MAX;
	if (!val->rv_primed) {
		val->rv_prev = cur;
		val->rv_prev_time = now;
		val->rv_primed = true;
		return;
	}
	/* Skip sub-ms re-samples; the next sample folds the full interval. */
	if (now - val->rv_prev_time < SBT_1MS)
		return;
	if (cur >= val->rv_prev)
		diff = cur - val->rv_prev;
	else
		diff = (UINT32_MAX - val->rv_prev) + cur + 1;
	val->rv_accum += diff;
	val->rv_prev = cur;
	val->rv_prev_time = now;
}

/* Sample one row's MSR on the current CPU and return the folded accumulator. */
static uint64_t
rapl_sample_row(int cpu, int ri)
{
	struct rapl_cpu *rc;
	uint64_t accum, cur;

	rc = rapl_pcpu[cpu];
	KASSERT(rc != NULL, ("[rapl,%d] null pcpu state cpu %d", __LINE__,
	    cpu));
	mtx_lock_spin(&rc->rc_mtx);
	if (rdmsr_safe(rapl_events[ri].re_msr, &cur) == 0)
		rapl_update_delta(&rc->rc_value[ri], cur);
	accum = rc->rc_value[ri].rv_accum;
	mtx_unlock_spin(&rc->rc_mtx);
	return (accum);
}

/* Guard rendezvous handler: sample every row on this CPU. */
static void
rapl_guard_handler(void *arg __unused)
{
	int cpu = curcpu;
	int ri;

	for (ri = 0; ri < rapl_npmcs; ri++)
		(void)rapl_sample_row(cpu, ri);
}

static void	rapl_guard_tick(void *arg);

/* (Re)arm the guard callout. Caller holds rapl_alloc_mtx. */
static void
rapl_guard_schedule(void)
{
	mtx_assert(&rapl_alloc_mtx, MA_OWNED);
	callout_reset_sbt(&rapl_guard_callout, rapl_guard_sbt,
	    rapl_guard_sbt / 10, rapl_guard_tick, NULL, 0);
}

/* Periodic overflow guard. */
static void
rapl_guard_tick(void *arg __unused)
{
	cpuset_t cpus;

	mtx_lock(&rapl_alloc_mtx);
	cpus = rapl_cpus;
	mtx_unlock(&rapl_alloc_mtx);

	if (!CPU_EMPTY(&cpus))
		smp_rendezvous_cpus(cpus, smp_no_rendezvous_barrier,
		    rapl_guard_handler, smp_no_rendezvous_barrier, NULL);

	/* Keep firing while any RAPL PMC remains allocated. */
	mtx_lock(&rapl_alloc_mtx);
	if (rapl_nalloc > 0)
		rapl_guard_schedule();
	mtx_unlock(&rapl_alloc_mtx);
}

static int
rapl_allocate_pmc(int cpu, int ri, struct pmc *pm __unused,
    const struct pmc_op_pmcallocate *a)
{

	KASSERT(cpu >= 0 && cpu < pmc_cpu_max(),
	    ("[rapl,%d] illegal CPU value %d", __LINE__, cpu));
	KASSERT(ri >= 0 && ri < rapl_npmcs,
	    ("[rapl,%d] illegal row index %d", __LINE__, ri));

	if (a->pm_class != PMC_CLASS_RAPL)
		return (EINVAL);

	if (a->pm_mode != PMC_MODE_SC)
		return (EINVAL);

	/* Power side channel (PLATYPUS): require privilege even if syspmcs bypass is set. */
	if (priv_check(curthread, PRIV_PMC_SYSTEM) != 0)
		return (EPERM);

	/* Reject events this vendor does not expose (e.g. DRAM on AMD). */
	if (a->pm_ev != rapl_events[ri].re_ev)
		return (EINVAL);

	/* Arm the guard on the first allocation (per-CPU and global). */
	mtx_lock(&rapl_alloc_mtx);
	if (rapl_pcpu[cpu]->rc_nalloc++ == 0)
		CPU_SET(cpu, &rapl_cpus);
	if (rapl_nalloc++ == 0)
		rapl_guard_schedule();
	mtx_unlock(&rapl_alloc_mtx);

	return (0);
}

static int
rapl_config_pmc(int cpu, int ri, struct pmc *pm)
{
	struct pmc_hw *phw;

	PMCDBG3(MDP,CFG,1, "cpu=%d ri=%d pm=%p", cpu, ri, pm);

	KASSERT(cpu >= 0 && cpu < pmc_cpu_max(),
	    ("[rapl,%d] illegal CPU value %d", __LINE__, cpu));
	KASSERT(ri >= 0 && ri < rapl_npmcs,
	    ("[rapl,%d] illegal row-index %d", __LINE__, ri));

	phw = &rapl_pcpu[cpu]->rc_hw[ri];

	KASSERT(pm == NULL || phw->phw_pmc == NULL,
	    ("[rapl,%d] pm=%p phw->pm=%p hwpmc not unconfigured", __LINE__,
	    pm, phw->phw_pmc));

	phw->phw_pmc = pm;

	return (0);
}

static int
rapl_describe(int cpu, int ri, struct pmc_info *pi, struct pmc **ppmc)
{
	const struct pmc_descr *pd;
	struct pmc_hw *phw;

	KASSERT(cpu >= 0 && cpu < pmc_cpu_max(),
	    ("[rapl,%d] illegal CPU %d", __LINE__, cpu));
	KASSERT(ri >= 0 && ri < rapl_npmcs,
	    ("[rapl,%d] illegal row-index %d", __LINE__, ri));

	phw = &rapl_pcpu[cpu]->rc_hw[ri];
	pd  = &rapl_pmcdesc[ri];

	strlcpy(pi->pm_name, pd->pd_name, sizeof(pi->pm_name));
	pi->pm_class = pd->pd_class;

	if (phw->phw_state & PMC_PHW_FLAG_IS_ENABLED) {
		pi->pm_enabled = true;
		*ppmc          = phw->phw_pmc;
	} else {
		pi->pm_enabled = false;
		*ppmc          = NULL;
	}

	return (0);
}

static int
rapl_get_config(int cpu, int ri, struct pmc **ppm)
{

	KASSERT(cpu >= 0 && cpu < pmc_cpu_max(),
	    ("[rapl,%d] illegal CPU %d", __LINE__, cpu));
	KASSERT(ri >= 0 && ri < rapl_npmcs,
	    ("[rapl,%d] illegal row-index %d", __LINE__, ri));

	*ppm = rapl_pcpu[cpu]->rc_hw[ri].phw_pmc;

	return (0);
}

static int
rapl_pcpu_init(struct pmc_mdep *md __unused, int cpu)
{
	struct pmc_cpu *pc;
	struct rapl_cpu *rapl_pc;
	int ri, n;

	KASSERT(cpu >= 0 && cpu < pmc_cpu_max(),
	    ("[rapl,%d] illegal cpu %d", __LINE__, cpu));
	KASSERT(rapl_pcpu, ("[rapl,%d] null pcpu", __LINE__));
	KASSERT(rapl_pcpu[cpu] == NULL, ("[rapl,%d] non-null per-cpu",
	    __LINE__));

	rapl_pc = malloc(sizeof(struct rapl_cpu), M_PMC, M_WAITOK | M_ZERO);
	mtx_init(&rapl_pc->rc_mtx, "rapl-cpu", NULL, MTX_SPIN);

	for (n = 0; n < rapl_npmcs; n++)
		rapl_pc->rc_hw[n].phw_state = PMC_PHW_FLAG_IS_ENABLED |
		    PMC_PHW_CPU_TO_STATE(cpu) | PMC_PHW_INDEX_TO_STATE(n) |
		    PMC_PHW_FLAG_IS_SHAREABLE;

	rapl_pcpu[cpu] = rapl_pc;

	KASSERT(pmc_pcpu, ("[rapl,%d] null generic pcpu", __LINE__));

	pc = pmc_pcpu[cpu];

	KASSERT(pc, ("[rapl,%d] null generic per-cpu", __LINE__));

	for (n = 0; n < rapl_npmcs; n++) {
		ri = rapl_ri + n;
		pc->pc_hwpmcs[ri] = &rapl_pc->rc_hw[n];
	}

	return (0);
}

static int
rapl_pcpu_fini(struct pmc_mdep *md __unused, int cpu)
{
	struct pmc_cpu *pc;
	int ri, n;

	KASSERT(cpu >= 0 && cpu < pmc_cpu_max(),
	    ("[rapl,%d] illegal cpu %d", __LINE__, cpu));
	KASSERT(rapl_pcpu[cpu] != NULL, ("[rapl,%d] null pcpu", __LINE__));
	KASSERT(rapl_pcpu[cpu]->rc_nalloc == 0,
	    ("[rapl,%d] %d PMCs still allocated on cpu %d", __LINE__,
	    rapl_pcpu[cpu]->rc_nalloc, cpu));

	/* Last release already drained the guard, so no handler can race here. */
	mtx_destroy(&rapl_pcpu[cpu]->rc_mtx);
	free(rapl_pcpu[cpu], M_PMC);
	rapl_pcpu[cpu] = NULL;

	pc = pmc_pcpu[cpu];
	for (n = 0; n < rapl_npmcs; n++) {
		ri = rapl_ri + n;
		pc->pc_hwpmcs[ri] = NULL;
	}

	return (0);
}

static int
rapl_read_pmc(int cpu, int ri, struct pmc *pm, pmc_value_t *v)
{
	enum pmc_mode mode __diagused;

	KASSERT(cpu >= 0 && cpu < pmc_cpu_max(),
	    ("[rapl,%d] illegal CPU value %d", __LINE__, cpu));
	KASSERT(ri >= 0 && ri < rapl_npmcs,
	    ("[rapl,%d] illegal ri %d", __LINE__, ri));

	mode = PMC_TO_MODE(pm);

	KASSERT(mode == PMC_MODE_SC,
	    ("[rapl,%d] illegal pmc mode %d", __LINE__, mode));

	PMCDBG1(MDP,REA,1, "rapl-read id=%d", ri);

	/* Bound to cpu by hwpmc, so rdmsr reads this CPU's domain (see DOMWIDE). */
	*v = rapl_raw_to_uj(rapl_sample_row(cpu, ri), rapl_events[ri].re_unit);

	return (0);
}

static int
rapl_release_pmc(int cpu, int ri, struct pmc *pmc __unused)
{
	struct pmc_hw *phw __diagused;
	bool last;

	KASSERT(cpu >= 0 && cpu < pmc_cpu_max(),
	    ("[rapl,%d] illegal CPU value %d", __LINE__, cpu));
	KASSERT(ri >= 0 && ri < rapl_npmcs,
	    ("[rapl,%d] illegal row-index %d", __LINE__, ri));

	phw = &rapl_pcpu[cpu]->rc_hw[ri];

	KASSERT(phw->phw_pmc == NULL,
	    ("[rapl,%d] PHW pmc %p non-NULL", __LINE__, phw->phw_pmc));

	mtx_lock(&rapl_alloc_mtx);
	KASSERT(rapl_pcpu[cpu]->rc_nalloc > 0 && rapl_nalloc > 0,
	    ("[rapl,%d] release underflow", __LINE__));
	if (--rapl_pcpu[cpu]->rc_nalloc == 0)
		CPU_CLR(cpu, &rapl_cpus);
	last = (--rapl_nalloc == 0);
	mtx_unlock(&rapl_alloc_mtx);

	/* Last release: drain the guard (sleepable here, mutex already dropped). */
	if (last)
		callout_drain(&rapl_guard_callout);

	return (0);
}

static int
rapl_start_pmc(int cpu __diagused, int ri __diagused, struct pmc *pm __unused)
{

	KASSERT(cpu >= 0 && cpu < pmc_cpu_max(),
	    ("[rapl,%d] illegal CPU value %d", __LINE__, cpu));
	KASSERT(ri >= 0 && ri < rapl_npmcs,
	    ("[rapl,%d] illegal row-index %d", __LINE__, ri));

	return (0);	/* RAPL counters are always running. */
}

static int
rapl_stop_pmc(int cpu __diagused, int ri __diagused, struct pmc *pm __unused)
{

	KASSERT(cpu >= 0 && cpu < pmc_cpu_max(),
	    ("[rapl,%d] illegal CPU value %d", __LINE__, cpu));
	KASSERT(ri >= 0 && ri < rapl_npmcs,
	    ("[rapl,%d] illegal row-index %d", __LINE__, ri));

	return (0);	/* RAPL counters cannot be stopped. */
}

static int
rapl_write_pmc(int cpu __diagused, int ri __diagused, struct pmc *pm __unused,
    pmc_value_t v __unused)
{

	KASSERT(cpu >= 0 && cpu < pmc_cpu_max(),
	    ("[rapl,%d] illegal CPU value %d", __LINE__, cpu));
	KASSERT(ri >= 0 && ri < rapl_npmcs,
	    ("[rapl,%d] illegal row-index %d", __LINE__, ri));

	/* Energy counters are not writable; refuse silently like TSC. */
	return (0);
}

/* Fault-safe RAPL MSR presence probe on the current CPU. */
static bool
rapl_msr_present(uint32_t msr)
{
	uint64_t v;

	return (rdmsr_safe(msr, &v) == 0);
}

/* Append an event row to the table if its MSR responds on this hardware. */
static void
rapl_add_event(enum pmc_event ev, uint32_t msr, uint32_t unit,
    const char *name)
{

	if (!rapl_msr_present(msr))
		return;

	rapl_events[rapl_npmcs].re_ev = ev;
	rapl_events[rapl_npmcs].re_msr = msr;
	rapl_events[rapl_npmcs].re_unit = unit;
	rapl_pmcdesc[rapl_npmcs].pd_class = PMC_CLASS_RAPL;
	rapl_pmcdesc[rapl_npmcs].pd_caps = RAPL_CAPS;
	rapl_pmcdesc[rapl_npmcs].pd_width = 64;
	strlcpy(rapl_pmcdesc[rapl_npmcs].pd_name, name,
	    sizeof(rapl_pmcdesc[rapl_npmcs].pd_name));
	rapl_npmcs++;
}

/* Guard interval: half the worst-case wrap period at RAPL_GUARD_WATT. */
static sbintime_t
rapl_compute_guard_sbt(uint32_t shift)
{
	uint64_t max_energy_uj, guard_ms;

	max_energy_uj = rapl_raw_to_uj(UINT32_MAX, shift);
	guard_ms = max_energy_uj / (2000ULL * RAPL_GUARD_WATT);
	if (guard_ms < RAPL_GUARD_MIN_MS)
		guard_ms = RAPL_GUARD_MIN_MS;
	else if (guard_ms > RAPL_GUARD_MAX_MS)
		guard_ms = RAPL_GUARD_MAX_MS;

	return (guard_ms * SBT_1MS);
}

/* Fixed 2^-16 J DRAM unit, not the ESU (HSX/KNL only -- not SPR/EMR/GNR). */
static bool
rapl_intel_fixed_dram_unit(void)
{

	if (CPUID_TO_FAMILY(cpu_id) != 0x6)
		return (false);

	switch (CPUID_TO_MODEL(cpu_id)) {
	case 0x3f:	/* Haswell-EP */
	case 0x4f:	/* Broadwell-EP */
	case 0x55:	/* Skylake/Cascade Lake/Cooper Lake-SP */
	case 0x56:	/* Broadwell-DE */
	case 0x57:	/* Xeon Phi KNL */
	case 0x6a:	/* Ice Lake-SP */
	case 0x6c:	/* Ice Lake-D */
	case 0x85:	/* Xeon Phi KNM */
		return (true);
	default:
		return (false);
	}
}

int
pmc_rapl_initialize(struct pmc_mdep *md, int maxcpu, int classindex)
{
	struct pmc_classdep *pcd;
	uint32_t unit_msr, pkg_msr, cores_msr, dram_msr;
	uint32_t esu, dram_unit, max_unit;
	uint64_t unit_val;
	int i;

	KASSERT(md != NULL, ("[rapl,%d] md is NULL", __LINE__));
	KASSERT(md->pmd_nclass >= 1, ("[rapl,%d] dubious md->nclass %d",
	    __LINE__, md->pmd_nclass));

	/* Select the per-vendor MSR set. */
	switch (cpu_vendor_id) {
	case CPU_VENDOR_AMD:
	case CPU_VENDOR_HYGON:
		unit_msr = MSR_AMD_RAPL_POWER_UNIT;
		pkg_msr = MSR_AMD_PKG_ENERGY_STATUS;
		cores_msr = MSR_AMD_CORE_ENERGY_STATUS;
		dram_msr = 0;			/* AMD has no DRAM domain */
		break;
	case CPU_VENDOR_INTEL:
		unit_msr = MSR_RAPL_POWER_UNIT;
		pkg_msr = MSR_PKG_ENERGY_STATUS;
		cores_msr = MSR_PP0_ENERGY_STATUS;
		dram_msr = MSR_DRAM_ENERGY_STATUS;
		break;
	default:
		return (ENXIO);
	}

	/* Decode the energy unit. */
	unit_val = rdmsr(unit_msr);
	esu = (unit_val >> 8) & 0x1f;
	dram_unit = rapl_intel_fixed_dram_unit() ? 16 : esu;

	/* Build the event table from the MSRs that actually respond. */
	rapl_npmcs = 0;
	rapl_add_event(PMC_EV_RAPL_ENERGY_PKG, pkg_msr, esu,
	    "RAPL_ENERGY_PKG");
	rapl_add_event(PMC_EV_RAPL_ENERGY_CORES, cores_msr, esu,
	    "RAPL_ENERGY_CORES");
	if (dram_msr != 0)
		rapl_add_event(PMC_EV_RAPL_ENERGY_DRAM, dram_msr, dram_unit,
		    "RAPL_ENERGY_DRAM");

	/* No RAPL energy MSR responded. */
	if (rapl_npmcs == 0)
		return (ENXIO);

	/* Size the guard for the fastest-wrapping row (largest unit shift). */
	max_unit = 0;
	for (i = 0; i < rapl_npmcs; i++)
		max_unit = MAX(max_unit, rapl_events[i].re_unit);
	rapl_guard_sbt = rapl_compute_guard_sbt(max_unit);
	rapl_nalloc = 0;
	CPU_ZERO(&rapl_cpus);

	mtx_init(&rapl_alloc_mtx, "rapl-alloc", NULL, MTX_DEF);
	/* It does not need associated mutex, the handler locks itself. */
	callout_init(&rapl_guard_callout, 1);

	rapl_pcpu = malloc(sizeof(struct rapl_cpu *) * maxcpu, M_PMC,
	    M_ZERO | M_WAITOK);

	pcd = &md->pmd_classdep[classindex];

	pcd->pcd_caps	= RAPL_CAPS;
	pcd->pcd_class	= PMC_CLASS_RAPL;
	pcd->pcd_num	= rapl_npmcs;
	pcd->pcd_ri	= md->pmd_npmc;
	pcd->pcd_width	= 64;

	pcd->pcd_allocate_pmc = rapl_allocate_pmc;
	pcd->pcd_config_pmc   = rapl_config_pmc;
	pcd->pcd_describe     = rapl_describe;
	pcd->pcd_get_config   = rapl_get_config;
	pcd->pcd_pcpu_init    = rapl_pcpu_init;
	pcd->pcd_pcpu_fini    = rapl_pcpu_fini;
	pcd->pcd_read_pmc     = rapl_read_pmc;
	pcd->pcd_release_pmc  = rapl_release_pmc;
	pcd->pcd_start_pmc    = rapl_start_pmc;
	pcd->pcd_stop_pmc     = rapl_stop_pmc;
	pcd->pcd_write_pmc    = rapl_write_pmc;

	rapl_ri = md->pmd_npmc;
	md->pmd_npmc += rapl_npmcs;

	return (0);
}

void
pmc_rapl_finalize(struct pmc_mdep *md __unused)
{
	PMCDBG0(MDP, INI, 1, "rapl-finalize");

	if (rapl_pcpu == NULL)
		return;

	KASSERT(rapl_nalloc == 0, ("[rapl,%d] %d PMCs still allocated",
	    __LINE__, rapl_nalloc));
	for (int i = 0; i < pmc_cpu_max(); i++)
		KASSERT(rapl_pcpu[i] == NULL, ("[rapl,%d] non-null pcpu cpu %d",
		    __LINE__, i));

	mtx_destroy(&rapl_alloc_mtx);

	free(rapl_pcpu, M_PMC);
	rapl_pcpu = NULL;
}