#!/usr/bin/perl

#  cudadl version 0.9: Compute discrete logs in smooth group orders
#                      using CUDA
#  Copyright (C) 2012 by Ryan Henry and Ian Goldberg
#                        {rhenry,iang}@cs.uwaterloo.ca
#
#  This program is free software: you can redistribute it and/or modify
#  it under the terms of version 3 of the GNU General Public License as
#  published by the Free Software Foundation.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see <http://www.gnu.org/licenses/>.

use strict;

die "Usage: $0 WORDS\n" unless $#ARGV == 0;

my $WORDS = shift;
my $WORDSm1 = ($WORDS-1);
my $WORDSp1 = ($WORDS+1);
my $WORDSp2 = ($WORDS+2);

my $i;

## Output an evil hack that declares a bunch of registers in the module
## scope

print <<EOA;
// An evil hack that declares a bunch of registers in the module scope
__device__ __noinline__ void evilhack() {
    asm("ret;}\\n"
        ".reg.u32 \$xr<$WORDSp1>;\\n"
        ".reg.u32 \$cg<$WORDS>;\\n"
        ".reg.u32 \$cy<$WORDS>;\\n"
	".reg.u32 \$yi,\$crp,\$m;\\n"
	".reg.u64 \$cs,\$ct;\\n"
	".reg.u32 \$zr<$WORDSp2>;\\n"
	".reg.u32 \$crho<$WORDS>;\\n"
	".func hackend {\\n"
EOA
print "\t\"ret.uni;\\n\"\n";

print "    );\n";
print "}\n";

print <<'EOA';
__device__ inline void loadx(unsigned int *x)
{
    asm(
EOA

for ($i=0;$i<$WORDS;++$i) {
    my $xoff = $i*4;
    print "\t\"ld.global.u32 \$xr$i, [%[x]+$xoff];\\n\"\n";
}
print "\t\"// Load c_rho into registers\\n\"\n";
my ($xoff, $xoffm1);
my $croff;
for($i=0;$i<$WORDS;++$i) {
    $croff = $i*4;
    print "\t\"ld.const.u32 \$crho$i, [c_rho+$croff];\\n\"\n";
}
print "\t\"mov.u32 \$crp, %[crp];\\n\"\n";
print "\t\"// Load c_g into registers\\n\"\n";
for($i=0;$i<$WORDS;++$i) {
    $croff = $i*4;
    print "\t\"ld.const.u32 \$cg$i, [c_g+$croff];\\n\"\n";
}
print "\t\"// Load c_y into registers\\n\"\n";
for($i=0;$i<$WORDS;++$i) {
    $croff = $i*4;
    print "\t\"ld.const.u32 \$cy$i, [c_y+$croff];\\n\"\n";
}


print <<'EOA';
    : : [x] "l" (x), [crp] "r" (c_rho_prime));
}

__device__ inline void savex(unsigned int *x)
{
    asm(
EOA

for ($i=0;$i<$WORDS;++$i) {
    my $xoff = $i*4;
    print "\t\"st.global.u32 [%[x]+$xoff], \$xr$i;\\n\"\n";
}

print <<'EOA';
    : : [x] "l" (x));
}
EOA

print <<'EOA';
__device__ void modmul(int multtype)
// x <- (x * y) * r_inv   (mod rho);
// where y is (multtype == 0 ? c_g : multtype == 1 ? c_y : x)
{
    asm(
    ".reg.pred $mtp0,$mtp1;\n"
EOA
print "    \".reg.pred \$retaddr<", $WORDS-2, ">;\\n\"\n";
print <<'EOA';
    "mov.u32 $m, 0;\n"
    "setp.eq.u32 $mtp0, %[tp], $m;\n"
    "mov.u32 $m, 1;\n"
    "setp.eq.u32 $mtp1, %[tp], $m;\n"
    : : [tp] "r" (multtype));
EOA

print "    asm(\n";
print "\t\"// First iteration: i=0 and z is unset\\n\"\n";

# Pick the correct value of $yi
print "\t\"selp.u32 \$yi, \$cg0, \$xr0, \$mtp0;\\n\"\n";
print "\t\"selp.u32 \$yi, \$cy0, \$yi, \$mtp1;\\n\"\n";

print "\t\"mul.wide.u32 \$cs, \$xr0, \$yi;\\n\"\n";
print "\t\"cvt.u32.u64 \$zr0, \$cs;\\n\"\n";
print "\t\"shr.u64 \$cs, \$cs, 32;\\n\"\n";

for($i=1; $i<$WORDS; ++$i) {
    $xoff = $i*32*4;
    print "\t\"mad.wide.u32 \$cs, \$xr$i, \$yi, \$cs;\\n\"\n";
    print "\t\"cvt.u32.u64 \$zr$i, \$cs;\\n\"\n";
    print "\t\"shr.u64 \$cs, \$cs, 32;\\n\"\n";
}
print "\t\"cvt.u32.u64 \$zr$WORDS, \$cs;\\n\"\n";

print "\t\"mul.lo.u32 \$m, \$zr0, \$crp;\\n\"\n";
print "\t\"cvt.u64.u32 \$ct, \$zr0;\\n\"\n";
print "\t\"mad.wide.u32 \$cs, \$m, \$crho0, \$ct;\\n\"\n";

for($i=1;$i<$WORDS;++$i) {
    my $im1 = ($i-1);
    print "\t\"shr.u64 \$cs, \$cs, 32;\\n\"\n";
    print "\t\"cvt.u64.u32 \$ct, \$zr$i;\\n\"\n";
    print "\t\"add.u64 \$cs, \$cs, \$ct;\\n\"\n";
    print "\t\"mad.wide.u32 \$cs, \$m, \$crho$i, \$cs;\\n\"\n";
    print "\t\"cvt.u32.u64 \$zr$im1, \$cs;\\n\"\n";
}
print "\t\"shr.u64 \$cs, \$cs, 32;\\n\"\n";
print "\t\"cvt.u64.u32 \$ct, \$zr$WORDS;\\n\"\n";
print "\t\"add.u64 \$cs, \$cs, \$ct;\\n\"\n";
print "\t\"cvt.u32.u64 \$zr$WORDSm1, \$cs;\\n\"\n";
print "\t\"shr.u64 \$cs, \$cs, 32;\\n\"\n";
print "\t\"cvt.u32.u64 \$zr$WORDS, \$cs;\\n\"\n";

my $j;
for ($j=1; $j<$WORDS-1; ++$j) {

print <<'EOA';
	"// Middle iterations\n"
EOA

# Pick the correct value of $yi
my $yioff = $j*4;
print "\t\"selp.u32 \$yi, \$cg$j, \$xr$j, \$mtp0;\\n\"\n";
print "\t\"selp.u32 \$yi, \$cy$j, \$yi, \$mtp1;\\n\"\n";
my $jm1 = $j-1;
my $jm2 = $j-2;
if ($j==1) {
    print "\t\"setp.eq.u32 \$retaddr0,\$xr0,\$xr0;\\n\"\n";
} else {
    print "\t\"setp.eq.u32 \$retaddr$jm1|\$retaddr$jm2,\$xr0,\$xr0;\\n\"\n";
}
print "\t\"bra.uni innerloop;\\n\"\n";
print "\t\"innerret$jm1:\\n\"\n";
}

print <<'EOA';
	"// Last iteration: store the output in x instead of z\n"
EOA

# Pick the correct value of $yi
my $yioff = ($WORDSm1)*4;
print "\t\"selp.u32 \$yi, \$cg$WORDSm1, \$xr$WORDSm1, \$mtp0;\\n\"\n";
print "\t\"selp.u32 \$yi, \$cy$WORDSm1, \$yi, \$mtp1;\\n\"\n";

print "\t\"cvt.u64.u32 \$ct, \$zr0;\\n\"\n";
print "\t\"mad.wide.u32 \$cs, \$xr0, \$yi, \$ct;\\n\"\n";
print "\t\"cvt.u32.u64 \$zr0, \$cs;\\n\"\n";
print "\t\"shr.u64 \$cs, \$cs, 32;\\n\"\n";

for($i=1; $i<$WORDS; ++$i) {
    $xoff = $i*32*4;
    print "\t\"cvt.u64.u32 \$ct, \$zr$i;\\n\"\n";
    print "\t\"add.u64 \$cs, \$cs, \$ct;\\n\"\n";
    print "\t\"mad.wide.u32 \$cs, \$xr$i, \$yi, \$cs;\\n\"\n";
    print "\t\"cvt.u32.u64 \$zr$i, \$cs;\\n\"\n";
    print "\t\"shr.u64 \$cs, \$cs, 32;\\n\"\n";
}
print "\t\"cvt.u64.u32 \$ct, \$zr$WORDS;\\n\"\n";
print "\t\"add.u64 \$cs, \$cs, \$ct;\\n\"\n";
print "\t\"cvt.u32.u64 \$zr$WORDS, \$cs;\\n\"\n";
print "\t\"shr.u64 \$cs, \$cs, 32;\\n\"\n";
print "\t\"cvt.u32.u64 \$zr$WORDSp1, \$cs;\\n\"\n";

print "\t\"mul.lo.u32 \$m, \$zr0, \$crp;\\n\"\n";
print "\t\"cvt.u64.u32 \$ct, \$zr0;\\n\"\n";
print "\t\"mad.wide.u32 \$cs, \$m, \$crho0, \$ct;\\n\"\n";

for($i=1;$i<$WORDS;++$i) {
    my $im1 = $i-1;
    print "\t\"shr.u64 \$cs, \$cs, 32;\\n\"\n";
    print "\t\"cvt.u64.u32 \$ct, \$zr$i;\\n\"\n";
    print "\t\"add.u64 \$cs, \$cs, \$ct;\\n\"\n";
    print "\t\"mad.wide.u32 \$cs, \$m, \$crho$i, \$cs;\\n\"\n";
    print "\t\"cvt.u32.u64 \$xr$im1, \$cs;\\n\"\n";
}
print "\t\"shr.u64 \$cs, \$cs, 32;\\n\"\n";
print "\t\"cvt.u64.u32 \$ct, \$zr$WORDS;\\n\"\n";
print "\t\"add.u64 \$cs, \$cs, \$ct;\\n\"\n";
print "\t\"cvt.u32.u64 \$xr$WORDSm1, \$cs;\\n\"\n";
print "\t\"shr.u64 \$cs, \$cs, 32;\\n\"\n";
print "\t\"cvt.u64.u32 \$ct, \$zr$WORDSp1;\\n\"\n";
print "\t\"add.u64 \$cs, \$cs, \$ct;\\n\"\n";
print "\t\"cvt.u32.u64 \$xr$WORDS, \$cs;\\n\"\n";

print <<'EOA';
        );
EOA

# If x >= c_rho, x <- x - c_rho

print "    asm(\n";
print "\t\".reg.pred \$gtP,\$ltP;\\n\"\n";
print "\t\"mov.u32 \$m, 0;\\n\"\n";
print "\t\"setp.gt.u32 \$gtP, \$xr$WORDS, \$m;\\n\"\n";
print "\t\"\@\$gtP bra \$comp_done;\\n\"\n";
for ($i = $WORDS-1; $i >= 1; --$i) {
    print "\t\"setp.gt.u32 \$gtP, \$xr$i, \$crho$i;\\n\"\n";
    print "\t\"\@\$gtP bra \$comp_done;\\n\"\n";
    print "\t\"setp.lt.u32 \$ltP, \$xr$i, \$crho$i;\\n\"\n";
    print "\t\"\@\$ltP bra \$comp_done;\\n\"\n";
}
print "\t\"setp.ge.u32 \$gtP, \$xr0, \$crho0;\\n\"\n";
print "\t\"\$comp_done:\\n\"\n";

# Do the subtraction
print "\t\"@\$gtP sub.cc.u32 \$xr0, \$xr0, \$crho0;\\n\"\n";
for ($i=1; $i<$WORDS; ++$i) {
    print "\t\"@\$gtP subc.cc.u32 \$xr$i, \$xr$i, \$crho$i;\\n\"\n";
}
print "\t\"@\$gtP subc.u32 \$xr$WORDS, \$xr$WORDS, \$m;\\n\"\n";

print "\t\"bra.uni mulmoddone;\\n\"\n";

print "\t\"innerloop:\\n\"\n";
print "\t\"cvt.u64.u32 \$ct, \$zr0;\\n\"\n";
print "\t\"mad.wide.u32 \$cs, \$xr0, \$yi, \$ct;\\n\"\n";
print "\t\"cvt.u32.u64 \$zr0, \$cs;\\n\"\n";
print "\t\"shr.u64 \$cs, \$cs, 32;\\n\"\n";

my $i;
for($i=1; $i<$WORDS; ++$i) {
    print "\t\"cvt.u64.u32 \$ct, \$zr$i;\\n\"\n";
    print "\t\"add.u64 \$cs, \$cs, \$ct;\\n\"\n";
    print "\t\"mad.wide.u32 \$cs, \$xr$i, \$yi, \$cs;\\n\"\n";
    print "\t\"cvt.u32.u64 \$zr$i, \$cs;\\n\"\n";
    print "\t\"shr.u64 \$cs, \$cs, 32;\\n\"\n";
}
print "\t\"cvt.u64.u32 \$ct, \$zr$WORDS;\\n\"\n";
print "\t\"add.u64 \$cs, \$cs, \$ct;\\n\"\n";
print "\t\"cvt.u32.u64 \$zr$WORDS, \$cs;\\n\"\n";
print "\t\"shr.u64 \$cs, \$cs, 32;\\n\"\n";
print "\t\"cvt.u32.u64 \$zr$WORDSp1, \$cs;\\n\"\n";

print "\t\"mul.lo.u32 \$m, \$zr0, \$crp;\\n\"\n";
print "\t\"cvt.u64.u32 \$ct, \$zr0;\\n\"\n";
print "\t\"mad.wide.u32 \$cs, \$m, \$crho0, \$ct;\\n\"\n";

for($i=1;$i<$WORDS;++$i) {
    my $im1 = ($i-1);
    print "\t\"shr.u64 \$cs, \$cs, 32;\\n\"\n";
    print "\t\"cvt.u64.u32 \$ct, \$zr$i;\\n\"\n";
    print "\t\"add.u64 \$cs, \$cs, \$ct;\\n\"\n";
    print "\t\"mad.wide.u32 \$cs, \$m, \$crho$i, \$cs;\\n\"\n";
    print "\t\"cvt.u32.u64 \$zr$im1, \$cs;\\n\"\n";
}
print "\t\"shr.u64 \$cs, \$cs, 32;\\n\"\n";
print "\t\"cvt.u64.u32 \$ct, \$zr$WORDS;\\n\"\n";
print "\t\"add.u64 \$cs, \$cs, \$ct;\\n\"\n";
print "\t\"cvt.u32.u64 \$zr$WORDSm1, \$cs;\\n\"\n";
print "\t\"shr.u64 \$cs, \$cs, 32;\\n\"\n";
print "\t\"cvt.u32.u64 \$zr$WORDS, \$cs;\\n\"\n";

for ($i=0;$i<$WORDS-2;++$i) {
    print "\t\"\@\$retaddr$i bra.uni innerret$i;\\n\"\n";
}
print "\t\"mulmoddone:\\n\"\n";

print "    );\n";
print "}\n";

print <<'EOA';

__device__ inline void write_dp(unsigned int *buf, unsigned int a_0,
    unsigned int a_1, unsigned int a_2, unsigned int b_0, unsigned int b_1,
    unsigned int b_2)
{
    unsigned short blockid = blockIdx.x + gridDim.x*blockIdx.y;
    unsigned short threadid = threadIdx.x + blockDim.x*threadIdx.y + blockDim.x*blockDim.y*threadIdx.z;
    buf[0] = (((unsigned int)blockid) << 16) + threadid;
    asm(
EOA
for($i=0;$i<$WORDS;++$i) {
    my $off = 4*($i+1);
    print "\t\"st.global.u32 [%0+$off], \$xr$i;\\n\"\n";
}
print <<'EOA';
    : : "l" (buf));
    buf[WORDS+1] = a_0;
    buf[WORDS+2] = a_1;
    buf[WORDS+3] = a_2;
    buf[WORDS+4] = b_0;
    buf[WORDS+5] = b_1;
    buf[WORDS+6] = b_2;
}
EOA