/**
 * Copyright (c) 2026 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
#include "aclnn_multinomial.h"
#include "stateless_sample_multinomial.h"
#include "random/stateless_exponential/op_api/stateless_exponential.h"
#include "multinomial_with_replacement.h"
#include "math/reduce_sum/op_api/reduce_sum_op.h"
#include "math/abs/op_api/abs.h"
#include "math/add/op_api/add.h"
#include "math/sub/op_api/sub.h"
#include "math/real_div/op_api/realdiv.h"
#include "math/log/op_api/log.h"
#include "math/arg_max_v2/op_api/argmax_v2.h"
#include "math/topk/op_api/topk.h"
#include "math/cumsum/op_api/cumsum.h"
#include "conversion/unsqueeze/op_host/op_api/unsqueeze.h"
#include "random/dsa_random_uniform/op_host/op_api/dsa_random_uniform.h"
#include "conversion/concat_d/op_api/concat_d.h"
#include "aclnn_kernels/cast.h"
#include "aclnn_kernels/reshape.h"
#include "math/greater_equal/op_api/greater_equal.h"
#include "conversion/view_copy/op_api/view_copy.h"
#include "aclnn_kernels/contiguous.h"
#include "aclnn_kernels/common/op_error_check.h"
#include "op_api/aclnn_check.h"
#include "opdev/common_types.h"
#include "opdev/shape_utils.h"
#include "opdev/op_dfx.h"
#include "opdev/op_executor.h"
#include "opdev/op_log.h"
#include "opdev/platform.h"

using namespace op;

#ifdef __cplusplus
extern "C" {
#endif

constexpr int64_t FLOAT32_MAX_CONSECUTIVE_INT = 1 << 24;
constexpr int64_t DIM_NUM_ONE = 1;
constexpr int64_t DIM_NUM_TWO = 2;
constexpr int64_t DIM_ZERO = 0;
constexpr int64_t CPU_NPU_BOUNDARY = 5000;

static const std::initializer_list<op::DataType> ASCEND910_DTYPE_SUPPORT_LIST = {
    op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_DOUBLE};

static const std::initializer_list<op::DataType> ASCEND910B_DTYPE_SUPPORT_LIST = {
    op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_DOUBLE, op::DataType::DT_BF16};

static const std::initializer_list<op::DataType> ASCEND950_DTYPE_SUPPORT_LIST = {
    op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_DOUBLE, op::DataType::DT_BF16};

static inline bool IsArch22(NpuArch arch) { return arch == NpuArch::DAV_2201; }

static inline bool IsArch35(NpuArch arch) { return arch == NpuArch::DAV_3510 || arch == NpuArch::DAV_9201; }

static bool UseAicpuPath(const aclTensor* self, int64_t selfSize)
{
    const auto curArch = GetCurrentPlatformInfo().GetCurNpuArch();
    if (IsArch35(curArch)) {
        const auto selfDtype = self->GetDataType();
        return selfDtype != op::DataType::DT_FLOAT && selfDtype != op::DataType::DT_FLOAT16 &&
               selfDtype != op::DataType::DT_BF16;
    }
    if (IsArch22(curArch)) {
        return selfSize <= CPU_NPU_BOUNDARY;
    }
    return true;
}

static bool ShouldCheckSeedOffsetDtype(int64_t selfSize)
{
    const auto curArch = GetCurrentPlatformInfo().GetCurNpuArch();
    if (IsArch22(curArch)) {
        return selfSize <= CPU_NPU_BOUNDARY;
    }
    return true;
}

static inline const std::initializer_list<op::DataType>& GetDtypeSupportList()
{
    const auto curArch = GetCurrentPlatformInfo().GetCurNpuArch();
    if (IsArch35(curArch)) {
        return ASCEND950_DTYPE_SUPPORT_LIST;
    }
    if (IsArch22(curArch)) {
        return ASCEND910B_DTYPE_SUPPORT_LIST;
    }
    return ASCEND910_DTYPE_SUPPORT_LIST;
}

static bool CheckNotNull(const aclTensor* self, const aclTensor* out)
{
    OP_CHECK_NULL(self, return false);
    OP_CHECK_NULL(out, return false);
    return true;
}

static bool CheckNotNullTensor(const aclTensor* self, const aclTensor* seedTensor, const aclTensor* offsetTensor,
                               const aclTensor* out)
{
    OP_CHECK_NULL(self, return false);
    OP_CHECK_NULL(seedTensor, return false);
    OP_CHECK_NULL(offsetTensor, return false);
    OP_CHECK_NULL(out, return false);
    return true;
}

static inline int64_t MakeWrapDim(int64_t dim, int64_t dimPostExpr)
{
    // support dim=0
    if (dimPostExpr <= 0) {
        dimPostExpr = 1;
    }
    if (dim < 0) {
        dim += dimPostExpr;
    }
    return dim;
}

static bool CheckDtypeValid(const aclTensor* self, const aclTensor* out)
{
    auto supportList = GetDtypeSupportList();
    OP_CHECK_DTYPE_NOT_SUPPORT(self, supportList, return false);
    OP_CHECK_DTYPE_NOT_SUPPORT(out, {op::DataType::DT_INT64}, return false);
    return true;
}

static bool CheckDtypeValidTensor(const aclTensor* self, const aclTensor* seedTensor, const aclTensor* offsetTensor,
                                  const aclTensor* out)
{
    auto supportList = GetDtypeSupportList();
    OP_CHECK_DTYPE_NOT_SUPPORT(self, supportList, return false);
    if (!(self->IsEmpty())) {
        const int64_t selfSize = self->GetViewShape().GetShapeSize();
        if (ShouldCheckSeedOffsetDtype(selfSize)) {
            OP_CHECK_DTYPE_NOT_SUPPORT(seedTensor, {op::DataType::DT_INT64}, return false);
            OP_CHECK_DTYPE_NOT_SUPPORT(offsetTensor, {op::DataType::DT_INT64}, return false);
        }
    }
    OP_CHECK_DTYPE_NOT_SUPPORT(out, {op::DataType::DT_INT64}, return false);
    return true;
}

static bool CheckShape(const aclTensor* self, int64_t numsamples, const aclTensor* out)
{
    if (self->GetViewShape().GetDimNum() != DIM_NUM_ONE && self->GetViewShape().GetDimNum() != DIM_NUM_TWO) {
        OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Dim of self can only be 1 or 2, but got %zu.",
                self->GetViewShape().GetDimNum());
        return false;
    }
    if (self->GetViewShape().GetDimNum() != out->GetViewShape().GetDimNum()) {
        OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Dim of self %zu should be equal to dim of out %zu.",
                self->GetViewShape().GetDimNum(), out->GetViewShape().GetDimNum());
        return false;
    }
    auto dimNum = out->GetViewShape().GetDimNum();
    auto nCategories = out->GetViewShape().GetDim(dimNum - 1);
    if (nCategories != numsamples) {
        OP_LOGE(ACLNN_ERR_PARAM_INVALID,
                "expected the size of out at last dim must be equal to numsamples %ld, but got %ld.", numsamples,
                nCategories);
        return false;
    }
    if (self->GetViewShape().GetDimNum() != DIM_NUM_ONE &&
        self->GetViewShape().GetDim(DIM_ZERO) != out->GetViewShape().GetDim(DIM_ZERO)) {
        OP_LOGE(ACLNN_ERR_PARAM_INVALID,
                "expected the size of out at first dim %ld must be equal to the size of self at first dim %ld when "
                "dimNum > 1",
                out->GetViewShape().GetDim(DIM_ZERO), self->GetViewShape().GetDim(DIM_ZERO));
        return false;
    }
    return true;
}

static bool CheckValueRange(const aclTensor* self, int64_t numsamples, bool replacement)
{
    if (numsamples <= 0) {
        OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Numsamples must be greater than 0, but got %ld.", numsamples);
        return false;
    }
    auto dimNum = self->GetViewShape().GetDimNum();
    auto nCategories = self->GetViewShape().GetDim(dimNum - 1);

    if (nCategories == 0) {
        OP_LOGE(ACLNN_ERR_PARAM_INVALID, "number of categories cannot be zero, got %ld.", nCategories);
        return false;
    }
    if (!replacement && (numsamples > nCategories)) {
        OP_LOGE(ACLNN_ERR_PARAM_INVALID,
                "numsamples %ld must be no more than the size of last dim %ld without replacement", numsamples,
                nCategories);
        return false;
    }
    if (nCategories > FLOAT32_MAX_CONSECUTIVE_INT) {
        OP_LOGE(ACLNN_ERR_PARAM_INVALID, "number of categories %ld cannot exceed 2^24", nCategories);
        return false;
    }
    return true;
}

static aclnnStatus CheckParams(const aclTensor* self, int64_t numsamples, bool replacement, const aclTensor* out)
{
    CHECK_RET(CheckNotNull(self, out), ACLNN_ERR_PARAM_NULLPTR);
    CHECK_RET(CheckDtypeValid(self, out), ACLNN_ERR_PARAM_INVALID);
    CHECK_RET(CheckShape(self, numsamples, out), ACLNN_ERR_PARAM_INVALID);
    CHECK_RET(CheckValueRange(self, numsamples, replacement), ACLNN_ERR_PARAM_INVALID);
    return ACLNN_SUCCESS;
}

static aclnnStatus CheckParamsTensor(const aclTensor* self, int64_t numsamples, bool replacement,
                                     const aclTensor* seedTensor, const aclTensor* offsetTensor, const aclTensor* out)
{
    CHECK_RET(CheckNotNullTensor(self, seedTensor, offsetTensor, out), ACLNN_ERR_PARAM_NULLPTR);
    CHECK_RET(CheckDtypeValidTensor(self, seedTensor, offsetTensor, out), ACLNN_ERR_PARAM_INVALID);
    CHECK_RET(CheckShape(self, numsamples, out), ACLNN_ERR_PARAM_INVALID);
    CHECK_RET(CheckValueRange(self, numsamples, replacement), ACLNN_ERR_PARAM_INVALID);
    return ACLNN_SUCCESS;
}

const aclTensor* GetRandomUniformReplaceMent(int64_t numsamples, const int64_t seed, const int64_t offset,
                                             aclOpExecutor* executor)
{
    const int64_t randAShape[] = {numsamples};
    auto randAShapeArray = executor->AllocIntArray(randAShape, 1);
    auto low = executor->AllocScalar(0.0f);
    auto high = executor->AllocScalar(1.0f);
    auto randomUniform = l0op::DSARandomUniform(randAShapeArray, seed, offset, low, high, executor);
    CHECK_RET(randomUniform != nullptr, nullptr);
    return randomUniform;
}

const aclTensor* GetRandomUniformReplaceMentTensor(int64_t numsamples, const aclTensor* seed, const aclTensor* offset,
                                                   aclOpExecutor* executor)
{
    // concat
    FVector<uint64_t> zeroVector = {static_cast<uint64_t>(0)};
    auto zeroTensor = executor->ConvertToTensor(zeroVector.data(), zeroVector.size(), offset->GetDataType());
    FVector<const aclTensor*> tensorListOnce;
    tensorListOnce.emplace_back(zeroTensor);
    tensorListOnce.emplace_back(offset);
    auto tensorList = executor->AllocTensorList(tensorListOnce.data(), tensorListOnce.size());
    auto concatTensor = l0op::ConcatD(tensorList, 0, executor);

    // RandomUniform, shape = {1, ..., 1, numsamples}
    const int64_t randAShape[] = {numsamples};
    auto randAShapeArray = executor->AllocIntArray(randAShape, 1);
    auto low = executor->AllocScalar(0.0f);
    auto high = executor->AllocScalar(1.0f);
    auto randomUniform = l0op::DSARandomUniformTensor(randAShapeArray, seed, concatTensor, low, high, executor);
    CHECK_RET(randomUniform != nullptr, nullptr);

    return randomUniform;
}

const aclTensor* RunMultinomialReplaceMent(const aclTensor* selfContiguous, int64_t numsamples,
                                           const aclTensor* randomUniform, const aclTensor* out,
                                           aclOpExecutor* executor)
{
    // trun weight into probability
    int64_t dimNum = static_cast<int64_t>(selfContiguous->GetViewShape().GetDimNum());
    int64_t lastDim = MakeWrapDim(-1, dimNum);
    const int64_t dim[] = {lastDim};
    auto dimArray = executor->AllocIntArray(dim, 1);
    auto sumSelf = l0op::ReduceSumOp(selfContiguous, dimArray, true, executor);
    CHECK_RET(sumSelf != nullptr, nullptr);

    auto divSumSelf = l0op::RealDiv(selfContiguous, sumSelf, executor);
    CHECK_RET(divSumSelf != nullptr, nullptr);

    // trun probability into point on [0, 1], shape = original_shape + {1}
    const aclTensor* dimTensor = nullptr;
    if (lastDim == 0 || lastDim > INT32_MAX) {
        dimTensor = executor->ConvertToTensor(&lastDim, 1, DataType::DT_INT64);
    } else {
        dimTensor = executor->ConvertToTensor(&lastDim, 1, DataType::DT_INT32);
    }
    CHECK_RET(dimTensor != nullptr, nullptr);
    auto accumSelf = l0op::Cumsum(divSumSelf, dimTensor, executor);
    CHECK_RET(accumSelf != nullptr, nullptr);

    auto unsqueezeAccum = l0op::UnsqueezeNd(accumSelf, -1, executor);
    CHECK_RET(unsqueezeAccum != nullptr, nullptr);

    auto castRandom = l0op::Cast(randomUniform, selfContiguous->GetDataType(), executor);
    CHECK_RET(castRandom != nullptr, nullptr);

    int64_t newShape[dimNum + 1];
    newShape[0] = numsamples;
    for (int64_t i = 0; i < dimNum; i++) {
        newShape[i] = 1;
    }
    newShape[dimNum] = numsamples;
    auto newShapeArray = executor->AllocIntArray(newShape, dimNum + 1);
    auto reshapeRandom = l0op::Reshape(castRandom, newShapeArray, executor);
    CHECK_RET(reshapeRandom != nullptr, nullptr);

    // caculate the point RandomUniform on interval unsqueezeAccum
    auto greaterEqual = l0op::GreaterEqual(reshapeRandom, unsqueezeAccum, executor);
    CHECK_RET(greaterEqual != nullptr, nullptr);

    auto castGreaterEqual = l0op::Cast(greaterEqual, out->GetDataType(), executor);
    CHECK_RET(castGreaterEqual != nullptr, nullptr);

    auto multinomialOut = l0op::ReduceSumOp(castGreaterEqual, dimArray, false, executor);
    CHECK_RET(multinomialOut != nullptr, nullptr);
    return multinomialOut;
}

static const aclTensor* Run950AicoreMultinomialWithReplacement(const aclTensor* selfContiguous, int64_t numsamples,
                                                               const aclTensor* seedTensor,
                                                               const aclTensor* offsetTensor, aclOpExecutor* executor)
{
    int64_t dimNum = static_cast<int64_t>(selfContiguous->GetViewShape().GetDimNum());
    int32_t lastDim = static_cast<int32_t>(MakeWrapDim(-1, dimNum));
    const int64_t dim[] = {lastDim};
    auto dimArray = executor->AllocIntArray(dim, 1);

    // Normalize weights to probabilities
    auto sumSelf = l0op::ReduceSumOp(selfContiguous, dimArray, true, executor);
    CHECK_RET(sumSelf != nullptr, nullptr);

    auto normProbsTensor = l0op::RealDiv(selfContiguous, sumSelf, executor);
    CHECK_RET(normProbsTensor != nullptr, nullptr);

    // Build CDF via cumulative sum
    const aclTensor* dimTensor = executor->ConvertToTensor(&lastDim, 1, DataType::DT_INT32);
    CHECK_RET(dimTensor != nullptr, nullptr);
    auto xTensor = l0op::Cumsum(normProbsTensor, dimTensor, executor);
    CHECK_RET(xTensor != nullptr, nullptr);

    // Launch fused RNG + binary search kernel, output is int64 indices directly
    auto multinomialOut = l0op::StatelessSampleMultinomial(xTensor, normProbsTensor, seedTensor, offsetTensor,
                                                           numsamples, executor);
    CHECK_RET(multinomialOut != nullptr, nullptr);
    return multinomialOut;
}

const aclTensor* GetRandomUniformNoReplaceMent(const aclTensor* selfContiguous, const int64_t seed,
                                               const int64_t offset, aclOpExecutor* executor)
{
    // exponentialOne = torch.exponential_(1)
    auto inputShape = op::ToShapeVector(selfContiguous->GetViewShape());
    auto inputShapeArray = executor->AllocIntArray(inputShape.data(), inputShape.size());
    CHECK_RET(inputShapeArray != nullptr, nullptr);
    auto low = executor->AllocScalar(0.0f);
    auto high = executor->AllocScalar(1.0f);
    auto randomUniform = l0op::DSARandomUniform(inputShapeArray, seed, offset, low, high, executor);
    CHECK_RET(randomUniform != nullptr, nullptr);

    return randomUniform;
}

const aclTensor* GetRandomUniformNoReplaceMentTensor(const aclTensor* selfContiguous, const aclTensor* seed,
                                                     const aclTensor* offset, aclOpExecutor* executor)
{
    // concat
    FVector<uint64_t> zeroVector = {static_cast<uint64_t>(0)};
    auto zeroTensor = executor->ConvertToTensor(zeroVector.data(), zeroVector.size(), offset->GetDataType());
    FVector<const aclTensor*> tensorListOnce;
    tensorListOnce.emplace_back(zeroTensor);
    tensorListOnce.emplace_back(offset);
    auto tensorList = executor->AllocTensorList(tensorListOnce.data(), tensorListOnce.size());
    auto concatTensor = l0op::ConcatD(tensorList, 0, executor);

    // exponentialOne = torch.exponential_(1)
    auto inputShape = op::ToShapeVector(selfContiguous->GetViewShape());
    auto inputShapeArray = executor->AllocIntArray(inputShape.data(), inputShape.size());
    CHECK_RET(inputShapeArray != nullptr, nullptr);
    auto low = executor->AllocScalar(0.0f);
    auto high = executor->AllocScalar(1.0f);
    auto randomUniform = l0op::DSARandomUniformTensor(inputShapeArray, seed, concatTensor, low, high, executor);
    CHECK_RET(randomUniform != nullptr, nullptr);

    return randomUniform;
}

static const aclTensor* Run950AicoreMultinomialWithoutReplacement(const aclTensor* selfContiguous, int64_t numsamples,
                                                                  const aclTensor* exp1Random, aclOpExecutor* executor)
{
    auto divExponential = l0op::RealDivInplace(selfContiguous, exp1Random, executor);
    CHECK_RET(divExponential != nullptr, nullptr);

    int64_t dimNum = static_cast<int64_t>(selfContiguous->GetViewShape().GetDimNum());
    int64_t lastDim = MakeWrapDim(-1, dimNum);

    const aclTensor* multinomialOut = nullptr;
    if (numsamples == 1) {
        auto argmaxOut = l0op::ArgMaxV2(divExponential, lastDim, true, op::DataType::DT_INT64, executor);
        CHECK_RET(argmaxOut != nullptr, nullptr);
        multinomialOut = l0op::Cast(argmaxOut, op::DataType::DT_INT64, executor);
    } else {
        auto topkOut = l0op::Topk(divExponential, numsamples, lastDim, true, true, op::DataType::DT_INT64, 1, executor);
        multinomialOut = l0op::Cast(std::get<1>(topkOut), op::DataType::DT_INT64, executor);
    }
    CHECK_RET(multinomialOut != nullptr, nullptr);
    return multinomialOut;
}

const aclTensor* RunMultinomialNoReplaceMent(const aclTensor* selfContiguous, int64_t numsamples,
                                             const aclTensor* randomUniform, const aclTensor* out,
                                             aclOpExecutor* executor)
{
    const float one = 1.0;
    const aclTensor* oneTensor = executor->ConvertToTensor(&one, 1, DataType::DT_FLOAT);
    auto subRandom = l0op::Sub(oneTensor, randomUniform, executor);
    CHECK_RET(subRandom != nullptr, nullptr);

    const float logBase = -1.0f;
    const float logScale = 1.0f;
    const float logShift = 0.0f;
    auto logRandom = l0op::Log(subRandom, logBase, logScale, logShift, executor);
    CHECK_RET(logRandom != nullptr, nullptr);

    auto absLog = l0op::Abs(logRandom, executor);
    CHECK_RET(absLog != nullptr, nullptr);

    // add positiveMin to prevent exponential = 0
    const float positiveMin = 1e-7;
    const aclTensor* positiveMinTensor = executor->ConvertToTensor(&positiveMin, 1, DataType::DT_FLOAT);
    auto exponential = l0op::Add(absLog, positiveMinTensor, executor);
    CHECK_RET(exponential != nullptr, nullptr);

    auto castExponential = l0op::Cast(exponential, selfContiguous->GetDataType(), executor);
    CHECK_RET(castExponential != nullptr, nullptr);

    // give selfContiguous random disturbance
    auto divExponential = l0op::RealDiv(selfContiguous, castExponential, executor);
    CHECK_RET(divExponential != nullptr, nullptr);

    int64_t dimNum = static_cast<int64_t>(selfContiguous->GetViewShape().GetDimNum());
    int64_t lastDim = MakeWrapDim(-1, dimNum);
    const aclTensor* multinomialOutInt32;
    if (numsamples == 1) {
        multinomialOutInt32 = l0op::ArgMaxV2(divExponential, lastDim, true, out->GetDataType(), executor);
    } else {
        auto topkOut = l0op::Topk(divExponential, numsamples, lastDim, true, true, op::DataType::DT_INT32, executor);
        multinomialOutInt32 = std::get<1>(topkOut);
    }
    CHECK_RET(multinomialOutInt32 != nullptr, nullptr);

    auto multinomialOut = l0op::Cast(multinomialOutInt32, out->GetDataType(), executor);
    CHECK_RET(multinomialOut != nullptr, nullptr);
    return multinomialOut;
}

aclnnStatus aclnnMultinomialGetWorkspaceSize(const aclTensor* self, int64_t numsamples, bool replacement, int64_t seed,
                                             int64_t offset, aclTensor* out, uint64_t* workspaceSize,
                                             aclOpExecutor** executor)
{
    OP_CHECK_COMM_INPUT(workspaceSize, executor);

    L2_DFX_PHASE_1(aclnnMultinomial, DFX_IN(self, numsamples, replacement, seed, offset), DFX_OUT(out));

    auto ret = CheckParams(self, numsamples, replacement, out);
    CHECK_RET(ret == ACLNN_SUCCESS, ret);
    auto uniqueExecutor = CREATE_EXECUTOR();
    CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR);

    if (self->IsEmpty()) {
        *workspaceSize = 0;
        uniqueExecutor.ReleaseTo(executor);
        return ACLNN_SUCCESS;
    }

    const int64_t dimNum = static_cast<int64_t>(self->GetViewShape().GetDimNum());
    int64_t selfSize = 1;
    for (int64_t i = 0; i < dimNum; i++) {
        selfSize *= self->GetViewShape().GetDim(i);
    }
    const auto curArch = GetCurrentPlatformInfo().GetCurNpuArch();

    auto selfContiguous = l0op::Contiguous(self, uniqueExecutor.get());
    CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_PARAM_NULLPTR);

    const aclTensor* multinomialOut;
    if (UseAicpuPath(self, selfSize)) {
        multinomialOut = l0op::MultinomialWithReplacement(selfContiguous, numsamples, replacement, seed, offset,
                                                          uniqueExecutor.get());
    } else if (!replacement || numsamples == 1) {
        if (IsArch35(curArch)) {
            // StatelessExponential takes tensor seed/offset; convert the scalar ones here.
            aclIntArray* seedList = uniqueExecutor->AllocIntArray(&seed, 1);
            CHECK_RET(seedList != nullptr, ACLNN_ERR_INNER_NULLPTR);
            auto seedTensor = uniqueExecutor->ConvertToTensor(seedList, op::DataType::DT_INT64);
            CHECK_RET(seedTensor != nullptr, ACLNN_ERR_INNER_NULLPTR);
            aclIntArray* offsetList = uniqueExecutor->AllocIntArray(&offset, 1);
            CHECK_RET(offsetList != nullptr, ACLNN_ERR_INNER_NULLPTR);
            auto offsetTensor = uniqueExecutor->ConvertToTensor(offsetList, op::DataType::DT_INT64);
            CHECK_RET(offsetTensor != nullptr, ACLNN_ERR_INNER_NULLPTR);
            // Pre-allocate a buffer filled in-place so the original weights (selfContiguous) are preserved.
            auto expInput = uniqueExecutor->AllocTensor(selfContiguous->GetViewShape(), selfContiguous->GetDataType(),
                                                        selfContiguous->GetViewFormat());
            CHECK_RET(expInput != nullptr, ACLNN_ERR_INNER_NULLPTR);
            auto exp1Random = l0op::StatelessExponential(expInput, seedTensor, offsetTensor, 1.0f,
                                                         uniqueExecutor.get());
            CHECK_RET(exp1Random != nullptr, ACLNN_ERR_PARAM_NULLPTR);
            multinomialOut = Run950AicoreMultinomialWithoutReplacement(selfContiguous, numsamples, exp1Random,
                                                                       uniqueExecutor.get());
        } else {
            auto randomUniform = GetRandomUniformNoReplaceMent(selfContiguous, seed, offset, uniqueExecutor.get());
            CHECK_RET(randomUniform != nullptr, ACLNN_ERR_PARAM_NULLPTR);
            multinomialOut = RunMultinomialNoReplaceMent(selfContiguous, numsamples, randomUniform, out,
                                                         uniqueExecutor.get());
        }
    } else {
        if (IsArch35(curArch)) {
            aclIntArray* seedList = uniqueExecutor->AllocIntArray(&seed, 1);
            CHECK_RET(seedList != nullptr, ACLNN_ERR_INNER_NULLPTR);
            auto seedTensor = uniqueExecutor->ConvertToTensor(seedList, op::DataType::DT_INT64);
            CHECK_RET(seedTensor != nullptr, ACLNN_ERR_INNER_NULLPTR);
            aclIntArray* offsetList = uniqueExecutor->AllocIntArray(&offset, 1);
            CHECK_RET(offsetList != nullptr, ACLNN_ERR_INNER_NULLPTR);
            auto offsetTensor = uniqueExecutor->ConvertToTensor(offsetList, op::DataType::DT_INT64);
            CHECK_RET(offsetTensor != nullptr, ACLNN_ERR_INNER_NULLPTR);
            multinomialOut = Run950AicoreMultinomialWithReplacement(selfContiguous, numsamples, seedTensor,
                                                                    offsetTensor, uniqueExecutor.get());
        } else {
            auto randomUniform = GetRandomUniformReplaceMent(numsamples, seed, offset, uniqueExecutor.get());
            CHECK_RET(randomUniform != nullptr, ACLNN_ERR_PARAM_NULLPTR);
            multinomialOut = RunMultinomialReplaceMent(selfContiguous, numsamples, randomUniform, out,
                                                       uniqueExecutor.get());
        }
    }
    CHECK_RET(multinomialOut != nullptr, ACLNN_ERR_PARAM_NULLPTR);

    auto viewCopyResult = l0op::ViewCopy(multinomialOut, out, uniqueExecutor.get());
    CHECK_RET(viewCopyResult != nullptr, ACLNN_ERR_PARAM_NULLPTR);

    *workspaceSize = uniqueExecutor->GetWorkspaceSize();
    uniqueExecutor.ReleaseTo(executor);

    return ACLNN_SUCCESS;
}

aclnnStatus aclnnMultinomial(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)
{
    L2_DFX_PHASE_2(aclnnMultinomial);
    return CommonOpExecutorRun(workspace, workspaceSize, executor, stream);
}

static const aclTensor* AddOffsetTensor(const aclTensor* offsetTensor, int64_t offset, aclOpExecutor* executor)
{
    FVector<int64_t> tmpVector = {static_cast<int64_t>(offset)};
    auto offsetTmpTensor = executor->ConvertToTensor(tmpVector.data(), tmpVector.size(), offsetTensor->GetDataType());
    CHECK_RET(offsetTmpTensor != nullptr, nullptr);
    auto offsetAddOut = l0op::Add(offsetTensor, offsetTmpTensor, executor);
    return offsetAddOut;
}

aclnnStatus aclnnMultinomialTensorGetWorkspaceSize(const aclTensor* self, int64_t numsamples, bool replacement,
                                                   const aclTensor* seedTensor, const aclTensor* offsetTensor,
                                                   int64_t offset, aclTensor* out, uint64_t* workspaceSize,
                                                   aclOpExecutor** executor)
{
    OP_CHECK_COMM_INPUT(workspaceSize, executor);

    L2_DFX_PHASE_1(aclnnMultinomialTensor, DFX_IN(self, numsamples, replacement, seedTensor, offsetTensor, offset),
                   DFX_OUT(out));

    auto ret = CheckParamsTensor(self, numsamples, replacement, seedTensor, offsetTensor, out);
    CHECK_RET(ret == ACLNN_SUCCESS, ret);
    auto uniqueExecutor = CREATE_EXECUTOR();
    CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR);

    if (self->IsEmpty()) {
        *workspaceSize = 0;
        uniqueExecutor.ReleaseTo(executor);
        return ACLNN_SUCCESS;
    }

    const int64_t dimNum = static_cast<int64_t>(self->GetViewShape().GetDimNum());
    int64_t selfSize = 1;
    for (int64_t i = 0; i < dimNum; i++) {
        selfSize *= self->GetViewShape().GetDim(i);
    }
    const auto curArch = GetCurrentPlatformInfo().GetCurNpuArch();

    auto selfContiguous = l0op::Contiguous(self, uniqueExecutor.get());
    CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_PARAM_NULLPTR);

    auto offsetAddOut = AddOffsetTensor(offsetTensor, offset, uniqueExecutor.get());
    CHECK_RET(offsetAddOut != nullptr, ACLNN_ERR_INNER_NULLPTR);

    const aclTensor* multinomialOut;
    if (UseAicpuPath(self, selfSize)) {
        multinomialOut = l0op::MultinomialWithReplacementTensor(selfContiguous, numsamples, replacement, seedTensor,
                                                                offsetAddOut, uniqueExecutor.get());
    } else if (!replacement || numsamples == 1) {
        if (IsArch35(curArch)) {
            auto expInput = uniqueExecutor->AllocTensor(selfContiguous->GetViewShape(), selfContiguous->GetDataType(),
                                                        selfContiguous->GetViewFormat());
            CHECK_RET(expInput != nullptr, ACLNN_ERR_INNER_NULLPTR);
            auto exp1Random = l0op::StatelessExponential(expInput, seedTensor, offsetAddOut, 1.0f,
                                                         uniqueExecutor.get());
            CHECK_RET(exp1Random != nullptr, ACLNN_ERR_PARAM_NULLPTR);
            multinomialOut = Run950AicoreMultinomialWithoutReplacement(selfContiguous, numsamples, exp1Random,
                                                                       uniqueExecutor.get());
        } else {
            auto randomUniform = GetRandomUniformNoReplaceMentTensor(selfContiguous, seedTensor, offsetAddOut,
                                                                     uniqueExecutor.get());
            CHECK_RET(randomUniform != nullptr, ACLNN_ERR_INNER_NULLPTR);
            multinomialOut = RunMultinomialNoReplaceMent(selfContiguous, numsamples, randomUniform, out,
                                                         uniqueExecutor.get());
        }
    } else {
        if (IsArch35(curArch)) {
            multinomialOut = Run950AicoreMultinomialWithReplacement(selfContiguous, numsamples, seedTensor,
                                                                    offsetAddOut, uniqueExecutor.get());
        } else {
            auto randomUniform = GetRandomUniformReplaceMentTensor(numsamples, seedTensor, offsetAddOut,
                                                                   uniqueExecutor.get());
            CHECK_RET(randomUniform != nullptr, ACLNN_ERR_INNER_NULLPTR);
            multinomialOut = RunMultinomialReplaceMent(selfContiguous, numsamples, randomUniform, out,
                                                       uniqueExecutor.get());
        }
    }
    CHECK_RET(multinomialOut != nullptr, ACLNN_ERR_PARAM_NULLPTR);

    auto viewCopyResult = l0op::ViewCopy(multinomialOut, out, uniqueExecutor.get());
    CHECK_RET(viewCopyResult != nullptr, ACLNN_ERR_PARAM_NULLPTR);

    *workspaceSize = uniqueExecutor->GetWorkspaceSize();
    uniqueExecutor.ReleaseTo(executor);

    return ACLNN_SUCCESS;
}

aclnnStatus aclnnMultinomialTensor(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)
{
    L2_DFX_PHASE_2(aclnnMultinomialTensor);
    return CommonOpExecutorRun(workspace, workspaceSize, executor, stream);
}

#ifdef __cplusplus
}
#endif