#!/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;
#define CIOS_REG_DECLARE \\
    asm volatile ( \\
        ".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" \\
    )

EOA

print <<'EOA';
#define CIOS_LOADX(x) \
    asm volatile ( \
EOA

for ($i=0;$i<$WORDS;++$i) {
    my $xoff = $i*4;
    print "\t\"ld.global.u32 \$xr$i, [%0+$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, %1;\\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';
    : : "l" (x), "r" (c_rho_prime))

EOA

print <<'EOA';
#define CIOS_SAVEX(x) \
    asm volatile ( \
EOA

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

print <<'EOA';
    : : "l" (x))

EOA

print <<'EOA';
// x <- (x * y) * r_inv   (mod rho);
// where y is (multtype == 0 ? c_g : multtype == 1 ? c_y : x)
#define CIOS_MODMUL(multtype) \
    asm volatile ( \
    ".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, %0, $m;\n" \
    "mov.u32 $m, 1;\n" \
    "setp.eq.u32 $mtp1, %0, $m;\n" \
    : : "r" (multtype)); \
EOA

print "    asm volatile ( \\\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 volatile ( \\\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 <<'EOA';

#define CIOS_WRITE_DP(buf, a_0, a_1, a_2, b_0, b_1, b_2) \
do { \
    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 volatile ( \
EOA
for($i=0;$i<$WORDS;++$i) {
    my $off = 4*($i+1);
    print "\t\"st.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; \
} while(0)
EOA