/**
 * aclnnDiv / aclnnInplaceDiv 整型除法精度复现（A2 910B2 / CANN 9.2.0）
 * 编译: g++ -O2 repro_div_int.cpp -o repro_div_int \
 *   -I${ASCEND_HOME_PATH}/include -L${ASCEND_HOME_PATH}/lib64 -lascendcl -lnnopbase -lopapi
 * 运行: ./repro_div_int [inplace]
 */
#include <cstdio>
#include <cstring>
#include <cstdint>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_div.h"

#define CK(x) do { auto r_ = (x); if (r_ != 0) { printf("[FAIL] %s rc=%d\n", #x, (int)r_); return 1; } } while (0)

static aclTensor* mk(const std::vector<int64_t>& shape, void* dev, aclDataType dt) {
    std::vector<int64_t> st(shape.size(), 1);
    for (int64_t i = (int64_t)shape.size() - 2; i >= 0; --i) st[i] = shape[i + 1] * st[i + 1];
    return aclCreateTensor(shape.data(), shape.size(), dt, st.data(), 0, ACL_FORMAT_ND,
                           shape.data(), shape.size(), dev);
}

static aclrtStream g_stream;

static bool run_div(aclDataType dt, void* d_s, void* d_o, void* d_out,
                    int64_t n, void* got, bool inplace) {
    std::vector<int64_t> shape = {n};
    aclTensor* ts = mk(shape, d_s, dt);
    aclTensor* to = mk(shape, d_o, dt);
    aclTensor* tout = mk(shape, d_out, dt);
    uint64_t ws = 0;
    aclOpExecutor* ex = nullptr;
    aclnnStatus st = inplace ? aclnnInplaceDivGetWorkspaceSize(ts, to, &ws, &ex)
                             : aclnnDivGetWorkspaceSize(ts, to, tout, &ws, &ex);
    if (st != 0) { printf("  第一段 rc=%d\n", (int)st); aclDestroyTensor(ts); aclDestroyTensor(to); aclDestroyTensor(tout); return false; }
    void* wsp = nullptr;
    if (ws > 0) CK(aclrtMalloc(&wsp, ws, ACL_MEM_MALLOC_HUGE_FIRST));
    st = inplace ? aclnnInplaceDiv(wsp, ws, ex, g_stream) : aclnnDiv(wsp, ws, ex, g_stream);
    aclrtSynchronizeStream(g_stream);
    aclrtMemcpy(got, n * (dt == ACL_INT64 ? 8 : 4), inplace ? d_s : d_out,
                n * (dt == ACL_INT64 ? 8 : 4), ACL_MEMCPY_DEVICE_TO_HOST);
    if (wsp) aclrtFree(wsp);
    aclDestroyTensor(ts); aclDestroyTensor(to); aclDestroyTensor(tout);
    return st == 0;
}

int main(int argc, char** argv) {
    bool inplace = argc > 1 && strcmp(argv[1], "inplace") == 0;
    CK(aclInit(nullptr));
    CK(aclrtSetDevice(0));
    CK(aclrtCreateStream(&g_stream));

    const int64_t N = 8;
    int32_t s32[N] = {2147483647, 2147483647, 1073741825, 2000000001,
                      -2147483647, 123456789, 1000000007, 16777217};
    int32_t o32[N] = {1, 2, 1, 3, 2, 7, 1000000007, 2};
    int64_t s64[N] = {4294967297LL, 1152921504606846977LL, 1000000000000000001LL,
                      1234567890123456789LL, 9007199254740993LL, 1000000000000001LL,
                      17179869184LL, -4611686018427387904LL};
    int64_t o64[N] = {1, 1, 3, 7, 1, 3, 1000000000LL, 2};

    void *ds = nullptr, *do_ = nullptr, *dout = nullptr;
    const size_t b32 = N * 4, b64 = N * 8;
    CK(aclrtMalloc(&ds, b64, ACL_MEM_MALLOC_HUGE_FIRST));
    CK(aclrtMalloc(&do_, b64, ACL_MEM_MALLOC_HUGE_FIRST));
    CK(aclrtMalloc(&dout, b64, ACL_MEM_MALLOC_HUGE_FIRST));

    printf("==== %s ====\n", inplace ? "aclnnInplaceDiv" : "aclnnDiv");
    int32_t got32[N] = {0};
    CK(aclrtMemcpy(ds, b32, s32, b32, ACL_MEMCPY_HOST_TO_DEVICE));
    CK(aclrtMemcpy(do_, b32, o32, b32, ACL_MEMCPY_HOST_TO_DEVICE));
    if (run_div(ACL_INT32, ds, do_, dout, N, got32, inplace)) {
        printf("INT32:\n");
        for (int i = 0; i < N; ++i) {
            int64_t exp = (int64_t)s32[i] / (int64_t)o32[i];
            printf("  %12d / %-12d 实测=%-12d 期望=%-12lld %s\n", s32[i], o32[i],
                   got32[i], (long long)exp, (int64_t)got32[i] == exp ? "OK" : "*** 错");
        }
    }
    int64_t got64[N] = {0};
    CK(aclrtMemcpy(ds, b64, s64, b64, ACL_MEMCPY_HOST_TO_DEVICE));
    CK(aclrtMemcpy(do_, b64, o64, b64, ACL_MEMCPY_HOST_TO_DEVICE));
    if (run_div(ACL_INT64, ds, do_, dout, N, got64, inplace)) {
        printf("INT64:\n");
        for (int i = 0; i < N; ++i) {
            int64_t exp = s64[i] / o64[i];
            printf("  %21lld / %-10lld 实测=%-21lld 期望=%-21lld %s\n", (long long)s64[i],
                   (long long)o64[i], (long long)got64[i], (long long)exp,
                   got64[i] == exp ? "OK" : "*** 错");
        }
    }
    aclrtFree(ds); aclrtFree(do_); aclrtFree(dout);
    aclrtDestroyStream(g_stream);
    aclrtResetDevice(0);
    aclFinalize();
    return 0;
}