zeknox/native/poseidon/element_bn128.cuh at main · okx/zeknox · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
// Copyright 2024 OKX Group
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

#ifndef __ZEKNOX_ELEMENT_BN128_CUH__
#define __ZEKNOX_ELEMENT_BN128_CUH__

#include "utils/cuda_utils.cuh"
#include "types/int_types.h"

#ifdef USE_CUDA
#define FFE FFElementGPU
#else
#define FFE FFElement
#endif

class FFE
{
private:
    DEVICE INLINE void set_vals(const u64 x[4])
    {
        z[0] = x[0];
        z[1] = x[1];
        z[2] = x[2];
        z[3] = x[3];
    }

    DEVICE INLINE void mul64(u64 a, u64 b, u64 *h, u64 *l);

    DEVICE INLINE void add64(u64 a, u64 b, u64 cin, u64 *r, u64 *cout);

    DEVICE INLINE void sub64(u64 a, u64 b, u64 bin, u64 *r, u64 *bout);

    DEVICE INLINE void madd0(u64 a, u64 b, u64 c, u64 *hi);

    DEVICE INLINE void madd1(u64 a, u64 b, u64 c, u64 *hi, u64 *lo);

    DEVICE INLINE void madd2(u64 a, u64 b, u64 c, u64 d, u64 *hi, u64 *lo);

    DEVICE INLINE void madd3(u64 a, u64 b, u64 c, u64 d, u64 e, u64 *hi, u64 *lo);

    DEVICE void mulGeneric(u64 *z, u64 *x, u64 *y);

    DEVICE void addGeneric(u64 *z, u64 *x, u64 *y);

    DEVICE void fromMontGeneric(u64 *z);

public:
    u64 z[4] = {0};

    DEVICE FFE(){};

    DEVICE FFE(const u64 x[4])
    {
        set_vals(x);
    };

    DEVICE static FFE NewElement()
    {
        return FFE();
    }

    DEVICE INLINE void Set(const FFE x)
    {
        set_vals(x.z);
    }

    // SetZero z = 0
    DEVICE INLINE void SetZero()
    {
        z[0] = 0;
        z[1] = 0;
        z[2] = 0;
        z[3] = 0;
    }

    // SetOne z = 1 (in Montgomery form)
    DEVICE INLINE void SetOne()
    {
        z[0] = 12436184717236109307u;
        z[1] = 3962172157175319849u;
        z[2] = 7381016538464732718u;
        z[3] = 1011752739694698287u;
    }

    DEVICE void SetUint64(u64 v);

    DEVICE void ToMont();

    DEVICE INLINE void FromMont()
    {
        fromMontGeneric(this->z);
    }

    // Add z = x + y mod q
    DEVICE INLINE void Add(FFE x, FFE y)
    {
        addGeneric(this->z, x.z, y.z);
    }

    // / Mul z = x * y mod q
    // see https://hackmd.io/@zkteam/modular_multiplication
    DEVICE INLINE void Mul(FFE x, FFE y)
    {
        mulGeneric(this->z, x.z, y.z);
    }

    DEVICE INLINE void Square(FFE x)
    {
        mulGeneric(this->z, x.z, x.z);
    }

    // Exp z = x^exponent mod q
    // Note: since only Exp(x,5) is used, we implement a custom version of Exp
    DEVICE INLINE void Exp5(const FFE x)
    {
        this->Set(x);
        this->Square(z);
        this->Square(z);
        u64 r[4];
        mulGeneric(r, z, (u64 *)x.z);
        this->set_vals(r);
    }

    DEVICE INLINE void Exp5()
    {
        u64 x[4];
        x[0] = z[0];
        x[1] = z[1];
        x[2] = z[2];
        x[3] = z[3];
        this->Square(z);
        this->Square(z);
        u64 r[4];
        mulGeneric(r, z, x);
        this->set_vals(r);
    }
};

#endif // __ZEKNOX_ELEMENT_BN128_CUH__