xref: /openssl/crypto/aes/asm/aes-riscv64-zkn.pl (revision 608cadfb) 
1#! /usr/bin/env perl
2# Copyright 2022 The OpenSSL Project Authors. All Rights Reserved.
3#
4# Licensed under the Apache License 2.0 (the "License").  You may not use
5# this file except in compliance with the License.  You can obtain a copy
6# in the file LICENSE in the source distribution or at
7# https://www.openssl.org/source/license.html
8
9# $output is the last argument if it looks like a file (it has an extension)
10# $flavour is the first argument if it doesn't look like a file
11$output = $#ARGV >= 0 && $ARGV[$#ARGV] =~ m|\.\w+$| ? pop : undef;
12$flavour = $#ARGV >= 0 && $ARGV[0] !~ m|\.| ? shift : undef;
13
14$output and open STDOUT,">$output";
15
16################################################################################
17# Utility functions to help with keeping track of which registers to stack/
18# unstack when entering / exiting routines.
19################################################################################
20{
21    # Callee-saved registers
22    my @callee_saved = map("x$_",(2,8,9,18..27));
23    # Caller-saved registers
24    my @caller_saved = map("x$_",(1,5..7,10..17,28..31));
25    my @must_save;
26    sub use_reg {
27        my $reg = shift;
28        if (grep(/^$reg$/, @callee_saved)) {
29            push(@must_save, $reg);
30        } elsif (!grep(/^$reg$/, @caller_saved)) {
31            # Register is not usable!
32            die("Unusable register ".$reg);
33        }
34        return $reg;
35    }
36    sub use_regs {
37        return map(use_reg("x$_"), @_);
38    }
39    sub save_regs {
40        my $ret = '';
41        my $stack_reservation = ($#must_save + 1) * 8;
42        my $stack_offset = $stack_reservation;
43        if ($stack_reservation % 16) {
44            $stack_reservation += 8;
45        }
46        $ret.="    addi    sp,sp,-$stack_reservation\n";
47        foreach (@must_save) {
48            $stack_offset -= 8;
49            $ret.="    sd      $_,$stack_offset(sp)\n";
50        }
51        return $ret;
52    }
53    sub load_regs {
54        my $ret = '';
55        my $stack_reservation = ($#must_save + 1) * 8;
56        my $stack_offset = $stack_reservation;
57        if ($stack_reservation % 16) {
58            $stack_reservation += 8;
59        }
60        foreach (@must_save) {
61            $stack_offset -= 8;
62            $ret.="    ld      $_,$stack_offset(sp)\n";
63        }
64        $ret.="    addi    sp,sp,$stack_reservation\n";
65        return $ret;
66    }
67    sub clear_regs {
68        @must_save = ();
69    }
70}
71
72################################################################################
73# util for encoding scalar crypto extension instructions
74################################################################################
75
76my @regs = map("x$_",(0..31));
77my %reglookup;
78@reglookup{@regs} = @regs;
79
80# Takes a register name, possibly an alias, and converts it to a register index
81# from 0 to 31
82sub read_reg {
83    my $reg = lc shift;
84    if (!exists($reglookup{$reg})) {
85        die("Unknown register ".$reg);
86    }
87    my $regstr = $reglookup{$reg};
88    if (!($regstr =~ /^x([0-9]+)$/)) {
89        die("Could not process register ".$reg);
90    }
91    return $1;
92}
93
94sub rv64_aes64ds {
95    # Encoding for aes64ds rd, rs1, rs2 instruction on RV64
96    #                XXXXXXX_ rs2 _ rs1 _XXX_ rd  _XXXXXXX
97    my $template = 0b0011101_00000_00000_000_00000_0110011;
98    my $rd = read_reg shift;
99    my $rs1 = read_reg shift;
100    my $rs2 = read_reg shift;
101
102    return ".word ".($template | ($rs2 << 20) | ($rs1 << 15) | ($rd << 7));
103}
104
105sub rv64_aes64dsm {
106    # Encoding for aes64dsm rd, rs1, rs2 instruction on RV64
107    #                XXXXXXX_ rs2 _ rs1 _XXX_ rd  _XXXXXXX
108    my $template = 0b0011111_00000_00000_000_00000_0110011;
109    my $rd = read_reg shift;
110    my $rs1 = read_reg shift;
111    my $rs2 = read_reg shift;
112
113    return ".word ".($template | ($rs2 << 20) | ($rs1 << 15) | ($rd << 7));
114}
115
116sub rv64_aes64es {
117    # Encoding for aes64es rd, rs1, rs2 instruction on RV64
118    #                XXXXXXX_ rs2 _ rs1 _XXX_ rd  _XXXXXXX
119    my $template = 0b0011001_00000_00000_000_00000_0110011;
120    my $rd = read_reg shift;
121    my $rs1 = read_reg shift;
122    my $rs2 = read_reg shift;
123
124    return ".word ".($template | ($rs2 << 20) | ($rs1 << 15) | ($rd << 7));
125}
126
127sub rv64_aes64esm {
128    # Encoding for aes64esm rd, rs1, rs2 instruction on RV64
129    #                XXXXXXX_ rs2 _ rs1 _XXX_ rd  _XXXXXXX
130    my $template = 0b0011011_00000_00000_000_00000_0110011;
131    my $rd = read_reg shift;
132    my $rs1 = read_reg shift;
133    my $rs2 = read_reg shift;
134
135    return ".word ".($template | ($rs2 << 20) | ($rs1 << 15) | ($rd << 7));
136}
137
138sub rv64_aes64im {
139    # Encoding for aes64im rd, rs1 instruction on RV64
140    #                XXXXXXXXXXXX_ rs1 _XXX_ rd  _XXXXXXX
141    my $template = 0b001100000000_00000_001_00000_0010011;
142    my $rd = read_reg shift;
143    my $rs1 = read_reg shift;
144
145    return ".word ".($template | ($rs1 << 15) | ($rd << 7));
146}
147
148sub rv64_aes64ks1i {
149    # Encoding for aes64ks1i rd, rs1, rnum instruction on RV64
150    #                XXXXXXXX_rnum_ rs1 _XXX_ rd  _XXXXXXX
151    my $template = 0b00110001_0000_00000_001_00000_0010011;
152    my $rd = read_reg shift;
153    my $rs1 = read_reg shift;
154    my $rnum = shift;
155
156    return ".word ".($template | ($rnum << 20) | ($rs1 << 15) | ($rd << 7));
157}
158
159sub rv64_aes64ks2 {
160    # Encoding for aes64ks2 rd, rs1, rs2 instruction on RV64
161    #                XXXXXXX_ rs2 _ rs1 _XXX_ rd  _XXXXXXX
162    my $template = 0b0111111_00000_00000_000_00000_0110011;
163    my $rd = read_reg shift;
164    my $rs1 = read_reg shift;
165    my $rs2 = read_reg shift;
166
167    return ".word ".($template | ($rs2 << 20) | ($rs1 << 15) | ($rd << 7));
168}
169################################################################################
170# Register assignment for rv64i_zkne_encrypt and rv64i_zknd_decrypt
171################################################################################
172
173# Registers to hold AES state (called s0-s3 or y0-y3 elsewhere)
174my ($Q0,$Q1,$Q2,$Q3) = use_regs(6..9);
175
176# Function arguments (x10-x12 are a0-a2 in the ABI)
177# Input block pointer, output block pointer, key pointer
178my ($INP,$OUTP,$KEYP) = use_regs(10..12);
179
180# Temporaries
181my ($T0,$T1) = use_regs(13..14);
182
183# Loop counter
184my ($loopcntr) = use_regs(30);
185
186################################################################################
187# void rv64i_zkne_encrypt(const unsigned char *in, unsigned char *out,
188#   const AES_KEY *key);
189################################################################################
190my $code .= <<___;
191.text
192.balign 16
193.globl rv64i_zkne_encrypt
194.type   rv64i_zkne_encrypt,\@function
195rv64i_zkne_encrypt:
196___
197
198$code .= save_regs();
199
200$code .= <<___;
201
202    # Load input to block cipher
203    ld      $Q0,0($INP)
204    ld      $Q1,8($INP)
205
206    # Load key
207    ld      $T0,0($KEYP)
208    ld      $T1,8($KEYP)
209
210    # Load number of rounds
211    lwu     $loopcntr,240($KEYP)
212
213    # initial transformation
214    xor     $Q0,$Q0,$T0
215    xor     $Q1,$Q1,$T1
216
217    # The main loop only executes the first N-1 rounds.
218    add     $loopcntr,$loopcntr,-1
219
220    # Do Nr - 1 rounds (final round is special)
2211:
222    @{[rv64_aes64esm $Q2,$Q0,$Q1]}
223    @{[rv64_aes64esm $Q3,$Q1,$Q0]}
224
225    # Update key ptr to point to next key in schedule
226    add     $KEYP,$KEYP,16
227
228    # Grab next key in schedule
229    ld      $T0,0($KEYP)
230    ld      $T1,8($KEYP)
231    xor     $Q0,$Q2,$T0
232    xor     $Q1,$Q3,$T1
233
234    add     $loopcntr,$loopcntr,-1
235    bgtz    $loopcntr,1b
236
237    # final round
238    @{[rv64_aes64es $Q2,$Q0,$Q1]}
239    @{[rv64_aes64es $Q3,$Q1,$Q0]}
240
241    # since not added 16 before
242    ld      $T0,16($KEYP)
243    ld      $T1,24($KEYP)
244    xor     $Q0,$Q2,$T0
245    xor     $Q1,$Q3,$T1
246
247    sd      $Q0,0($OUTP)
248    sd      $Q1,8($OUTP)
249
250    # Pop registers and return
251___
252
253$code .= load_regs();
254
255$code .= <<___;
256    ret
257___
258
259################################################################################
260# void rv64i_zknd_decrypt(const unsigned char *in, unsigned char *out,
261#   const AES_KEY *key);
262################################################################################
263$code .= <<___;
264.text
265.balign 16
266.globl rv64i_zknd_decrypt
267.type   rv64i_zknd_decrypt,\@function
268rv64i_zknd_decrypt:
269___
270
271$code .= save_regs();
272
273$code .= <<___;
274
275    # Load input to block cipher
276    ld      $Q0,0($INP)
277    ld      $Q1,8($INP)
278
279    # Load number of rounds
280    lwu     $loopcntr,240($KEYP)
281
282    # Load the last key
283    slli    $T0,$loopcntr,4
284    add     $KEYP,$KEYP,$T0
285    ld      $T0,0($KEYP)
286    ld      $T1,8($KEYP)
287
288    xor     $Q0,$Q0,$T0
289    xor     $Q1,$Q1,$T1
290
291    # The main loop only executes the first N-1 rounds.
292    add     $loopcntr,$loopcntr,-1
293
294    # Do Nr - 1 rounds (final round is special)
2951:
296    @{[rv64_aes64dsm $Q2,$Q0,$Q1]}
297    @{[rv64_aes64dsm $Q3,$Q1,$Q0]}
298
299    # Update key ptr to point to next key in schedule
300    add     $KEYP,$KEYP,-16
301
302    # Grab next key in schedule
303    ld      $T0,0($KEYP)
304    ld      $T1,8($KEYP)
305    xor     $Q0,$Q2,$T0
306    xor     $Q1,$Q3,$T1
307
308    add     $loopcntr,$loopcntr,-1
309    bgtz    $loopcntr,1b
310
311    # final round
312    @{[rv64_aes64ds $Q2,$Q0,$Q1]}
313    @{[rv64_aes64ds $Q3,$Q1,$Q0]}
314
315    add     $KEYP,$KEYP,-16
316    ld      $T0,0($KEYP)
317    ld      $T1,8($KEYP)
318    xor     $Q0,$Q2,$T0
319    xor     $Q1,$Q3,$T1
320
321    sd      $Q0,0($OUTP)
322    sd      $Q1,8($OUTP)
323    # Pop registers and return
324___
325
326$code .= load_regs();
327
328$code .= <<___;
329    ret
330___
331
332clear_regs();
333
334################################################################################
335# Register assignment for rv64i_zkn[e/d]_set_[en/de]crypt_key
336################################################################################
337
338# Function arguments (x10-x12 are a0-a2 in the ABI)
339# Pointer to user key, number of bits in key, key pointer
340my ($UKEY,$BITS,$KEYP) = use_regs(10..12);
341
342# Temporaries
343my ($T0,$T1,$T2,$T3,$T4) = use_regs(6..8,13..14);
344
345################################################################################
346# utility functions for rv64i_zkne_set_encrypt_key
347################################################################################
348sub ke128enc {
349    my $rnum = 0;
350    my $ret = '';
351$ret .= <<___;
352    ld      $T0,0($UKEY)
353    ld      $T1,8($UKEY)
354    sd      $T0,0($KEYP)
355    sd      $T1,8($KEYP)
356___
357    while($rnum < 10) {
358$ret .= <<___;
359    @{[rv64_aes64ks1i   $T2,$T1,$rnum]}
360    @{[rv64_aes64ks2    $T0,$T2,$T0]}
361    @{[rv64_aes64ks2    $T1,$T0,$T1]}
362    add         $KEYP,$KEYP,16
363    sd          $T0,0($KEYP)
364    sd          $T1,8($KEYP)
365___
366        $rnum++;
367    }
368    return $ret;
369}
370
371sub ke192enc {
372    my $rnum = 0;
373    my $ret = '';
374$ret .= <<___;
375    ld      $T0,0($UKEY)
376    ld      $T1,8($UKEY)
377    ld      $T2,16($UKEY)
378    sd      $T0,0($KEYP)
379    sd      $T1,8($KEYP)
380    sd      $T2,16($KEYP)
381___
382    while($rnum < 8) {
383$ret .= <<___;
384    @{[rv64_aes64ks1i   $T3,$T2,$rnum]}
385    @{[rv64_aes64ks2    $T0,$T3,$T0]}
386    @{[rv64_aes64ks2    $T1,$T0,$T1]}
387___
388        if ($rnum != 7) {
389        # note that (8+1)*24 = 216, (12+1)*16 = 208
390        # thus the last 8 bytes can be dropped
391$ret .= <<___;
392    @{[rv64_aes64ks2    $T2,$T1,$T2]}
393___
394        }
395$ret .= <<___;
396    add         $KEYP,$KEYP,24
397    sd          $T0,0($KEYP)
398    sd          $T1,8($KEYP)
399___
400        if ($rnum != 7) {
401$ret .= <<___;
402    sd          $T2,16($KEYP)
403___
404        }
405        $rnum++;
406    }
407    return $ret;
408}
409
410sub ke256enc {
411    my $rnum = 0;
412    my $ret = '';
413$ret .= <<___;
414    ld      $T0,0($UKEY)
415    ld      $T1,8($UKEY)
416    ld      $T2,16($UKEY)
417    ld      $T3,24($UKEY)
418    sd      $T0,0($KEYP)
419    sd      $T1,8($KEYP)
420    sd      $T2,16($KEYP)
421    sd      $T3,24($KEYP)
422___
423    while($rnum < 7) {
424$ret .= <<___;
425    @{[rv64_aes64ks1i   $T4,$T3,$rnum]}
426    @{[rv64_aes64ks2    $T0,$T4,$T0]}
427    @{[rv64_aes64ks2    $T1,$T0,$T1]}
428    add         $KEYP,$KEYP,32
429    sd          $T0,0($KEYP)
430    sd          $T1,8($KEYP)
431___
432        if ($rnum != 6) {
433        # note that (7+1)*32 = 256, (14+1)*16 = 240
434        # thus the last 16 bytes can be dropped
435$ret .= <<___;
436    @{[rv64_aes64ks1i   $T4,$T1,0xA]}
437    @{[rv64_aes64ks2    $T2,$T4,$T2]}
438    @{[rv64_aes64ks2    $T3,$T2,$T3]}
439    sd          $T2,16($KEYP)
440    sd          $T3,24($KEYP)
441___
442        }
443        $rnum++;
444    }
445    return $ret;
446}
447
448################################################################################
449# void rv64i_zkne_set_encrypt_key(const unsigned char *userKey, const int bits,
450#   AES_KEY *key)
451################################################################################
452sub AES_set_common {
453    my ($ke128, $ke192, $ke256) = @_;
454    my $ret = '';
455$ret .= <<___;
456    bnez    $UKEY,1f        # if (!userKey || !key) return -1;
457    bnez    $KEYP,1f
458    li      a0,-1
459    ret
4601:
461    # Determine number of rounds from key size in bits
462    li      $T0,128
463    bne     $BITS,$T0,1f
464    li      $T1,10          # key->rounds = 10 if bits == 128
465    sw      $T1,240($KEYP)  # store key->rounds
466$ke128
467    j       4f
4681:
469    li      $T0,192
470    bne     $BITS,$T0,2f
471    li      $T1,12          # key->rounds = 12 if bits == 192
472    sw      $T1,240($KEYP)  # store key->rounds
473$ke192
474    j       4f
4752:
476    li      $T1,14          # key->rounds = 14 if bits == 256
477    li      $T0,256
478    beq     $BITS,$T0,3f
479    li      a0,-2           # If bits != 128, 192, or 256, return -2
480    j       5f
4813:
482    sw      $T1,240($KEYP)  # store key->rounds
483$ke256
4844:  # return 0
485    li      a0,0
4865:  # return a0
487___
488    return $ret;
489}
490$code .= <<___;
491.text
492.balign 16
493.globl rv64i_zkne_set_encrypt_key
494.type   rv64i_zkne_set_encrypt_key,\@function
495rv64i_zkne_set_encrypt_key:
496___
497$code .= save_regs();
498$code .= AES_set_common(ke128enc(), ke192enc(),ke256enc());
499$code .= load_regs();
500$code .= <<___;
501    ret
502___
503
504################################################################################
505# utility functions for rv64i_zknd_set_decrypt_key
506################################################################################
507sub ke128dec {
508    my $rnum = 0;
509    my $ret = '';
510$ret .= <<___;
511    ld      $T0,0($UKEY)
512    ld      $T1,8($UKEY)
513    sd      $T0,0($KEYP)
514    sd      $T1,8($KEYP)
515___
516    while($rnum < 10) {
517$ret .= <<___;
518    @{[rv64_aes64ks1i   $T2,$T1,$rnum]}
519    @{[rv64_aes64ks2    $T0,$T2,$T0]}
520    @{[rv64_aes64ks2    $T1,$T0,$T1]}
521    add         $KEYP,$KEYP,16
522___
523    # need to aes64im for [1:N-1] round keys
524    # this is from the fact that aes64dsm subwords first then mix column
525    # intuitively decryption needs to first mix column then subwords
526    # however, for merging datapaths (encryption first subwords then mix column)
527    # aes64dsm chooses to inverse the order of them, thus
528    # transform should then be done on the round key
529        if ($rnum < 9) {
530$ret .= <<___;
531    @{[rv64_aes64im     $T2,$T0]}
532    sd          $T2,0($KEYP)
533    @{[rv64_aes64im     $T2,$T1]}
534    sd          $T2,8($KEYP)
535___
536        } else {
537$ret .= <<___;
538    sd          $T0,0($KEYP)
539    sd          $T1,8($KEYP)
540___
541        }
542        $rnum++;
543    }
544    return $ret;
545}
546
547sub ke192dec {
548    my $rnum = 0;
549    my $ret = '';
550$ret .= <<___;
551    ld      $T0,0($UKEY)
552    ld      $T1,8($UKEY)
553    ld      $T2,16($UKEY)
554    sd      $T0,0($KEYP)
555    sd      $T1,8($KEYP)
556    @{[rv64_aes64im $T3,$T2]}
557    sd      $T3,16($KEYP)
558___
559    while($rnum < 8) {
560$ret .= <<___;
561    @{[rv64_aes64ks1i   $T3,$T2,$rnum]}
562    @{[rv64_aes64ks2    $T0,$T3,$T0]}
563    @{[rv64_aes64ks2    $T1,$T0,$T1]}
564    add         $KEYP,$KEYP,24
565___
566        if ($rnum < 7) {
567$ret .= <<___;
568    @{[rv64_aes64im     $T3,$T0]}
569    sd          $T3,0($KEYP)
570    @{[rv64_aes64im     $T3,$T1]}
571    sd          $T3,8($KEYP)
572    # the reason is in ke192enc
573    @{[rv64_aes64ks2    $T2,$T1,$T2]}
574    @{[rv64_aes64im     $T3,$T2]}
575    sd          $T3,16($KEYP)
576___
577        } else { # rnum == 7
578$ret .= <<___;
579    sd          $T0,0($KEYP)
580    sd          $T1,8($KEYP)
581___
582        }
583        $rnum++;
584    }
585    return $ret;
586}
587
588sub ke256dec {
589    my $rnum = 0;
590    my $ret = '';
591$ret .= <<___;
592    ld      $T0,0($UKEY)
593    ld      $T1,8($UKEY)
594    ld      $T2,16($UKEY)
595    ld      $T3,24($UKEY)
596    sd      $T0,0($KEYP)
597    sd      $T1,8($KEYP)
598    @{[rv64_aes64im $T4,$T2]}
599    sd      $T4,16($KEYP)
600    @{[rv64_aes64im $T4,$T3]}
601    sd      $T4,24($KEYP)
602___
603    while($rnum < 7) {
604$ret .= <<___;
605    @{[rv64_aes64ks1i   $T4,$T3,$rnum]}
606    @{[rv64_aes64ks2    $T0,$T4,$T0]}
607    @{[rv64_aes64ks2    $T1,$T0,$T1]}
608    add         $KEYP,$KEYP,32
609___
610        if ($rnum < 6) {
611$ret .= <<___;
612    @{[rv64_aes64ks1i   $T4,$T1,0xA]}
613    @{[rv64_aes64ks2    $T2,$T4,$T2]}
614    @{[rv64_aes64ks2    $T3,$T2,$T3]}
615    @{[rv64_aes64im     $T4,$T0]}
616    sd          $T4,0($KEYP)
617    @{[rv64_aes64im     $T4,$T1]}
618    sd          $T4,8($KEYP)
619    @{[rv64_aes64im     $T4,$T2]}
620    sd          $T4,16($KEYP)
621    @{[rv64_aes64im     $T4,$T3]}
622    sd          $T4,24($KEYP)
623___
624        } else {
625$ret .= <<___;
626    sd          $T0,0($KEYP)
627    sd          $T1,8($KEYP)
628    # last two one dropped
629___
630        }
631        $rnum++;
632    }
633    return $ret;
634}
635
636################################################################################
637# void rv64i_zknd_set_decrypt_key(const unsigned char *userKey, const int bits,
638#   AES_KEY *key)
639################################################################################
640$code .= <<___;
641.text
642.balign 16
643.globl rv64i_zknd_set_decrypt_key
644.type   rv64i_zknd_set_decrypt_key,\@function
645rv64i_zknd_set_decrypt_key:
646___
647$code .= save_regs();
648$code .= AES_set_common(ke128dec(), ke192dec(),ke256dec());
649$code .= load_regs();
650$code .= <<___;
651    ret
652___
653
654print $code;
655close STDOUT or die "error closing STDOUT: $!";
656