<s> '''<STR_LIT>''' <EOL> import numpy as np <EOL> import time <EOL> from numba import * <EOL> @ cuda . jit ( argtypes = [ f8 [ : , : ] , f8 [ : , : ] , f8 [ : , : ] ] ) <EOL> def jacobi_relax_core ( A , Anew , error ) : <EOL> smem = cuda . shared . array ( shape = ( <NUM_LIT:32> + <NUM_LIT:2> , <NUM_LIT:32> + <NUM_LIT:2> ) , dtype = f4 ) <EOL> n = A . shape [ <NUM_LIT:0> ] <EOL> m = A . shape [ <NUM_LIT:1> ] <EOL> tx = cuda . threadIdx . x <EOL> ty = cuda . threadIdx . y <EOL> j = ty + cuda . blockIdx . y * cuda . blockDim . y <EOL> i = tx + cuda . blockIdx . x * cuda . blockDim . x <EOL> sy = ty + <NUM_LIT:1> <EOL> sx = tx + <NUM_LIT:1> <EOL> smem [ sy , sx ] = A [ j , i ] <EOL> if tx == <NUM_LIT:0> and i >= <NUM_LIT:1> : <EOL> smem [ sy , <NUM_LIT:0> ] = A [ j , i - <NUM_LIT:1> ] <EOL> if ty == <NUM_LIT:0> and j < m - <NUM_LIT:1> : <EOL> smem [ <NUM_LIT:0> , sx ] = A [ j - <NUM_LIT:1> , i ] <EOL> if tx == <NUM_LIT> and j >= <NUM_LIT:1> : <EOL> smem [ sy , <NUM_LIT> ] = A [ j , i + <NUM_LIT:1> ] <EOL> if ty == <NUM_LIT> and j < n - <NUM_LIT:1> : <EOL> smem [ <NUM_LIT> , sx ] = A [ j + <NUM_LIT:1> , i ] <EOL> cuda . syncthreads ( ) <EOL> if j >= <NUM_LIT:1> and j < n - <NUM_LIT:1> and i >= <NUM_LIT:1> and i < m - <NUM_LIT:1> : <EOL> Anew [ j , i ] = <NUM_LIT> * ( smem [ sy , sx + <NUM_LIT:1> ] + smem [ sy , sx - <NUM_LIT:1> ] + smem [ sy - <NUM_LIT:1> , sx ] + smem [ sy + <NUM_LIT:1> , sx ] ) <EOL> error [ j , i ] = Anew [ j , i ] - A [ j , i ] <EOL> def main ( ) : <EOL> NN = <NUM_LIT> <EOL> NM = <NUM_LIT> <EOL> A = np . zeros ( ( NN , NM ) , dtype = np . float64 ) <EOL> Anew = np . zeros ( ( NN , NM ) , dtype = np . float64 ) <EOL> n = NN <EOL> m = NM <EOL> iter_max = <NUM_LIT:1000> <EOL> tol = <NUM_LIT> <EOL> error = <NUM_LIT:1.0> <EOL> for j in range ( n ) : <EOL> A [ j , <NUM_LIT:0> ] = <NUM_LIT:1.0> <EOL> Anew [ j , <NUM_LIT:0> ] = <NUM_LIT:1.0> <EOL> print ( "<STR_LIT>" % ( n , m ) ) <EOL> timer = time . time ( ) <EOL> iter = <NUM_LIT:0> <EOL> blockdim = ( <NUM_LIT:32> , <NUM_LIT:32> ) <EOL> griddim = ( NN // blockdim [ <NUM_LIT:0> ] , NM // blockdim [ <NUM_LIT:1> ] ) <EOL> error_grid = np . zeros_like ( A ) <EOL> stream = cuda . stream ( ) <EOL> dA = cuda . to_device ( A , stream ) <EOL> dAnew = cuda . to_device ( Anew , stream ) <EOL> derror_grid = cuda . to_device ( error_grid , stream ) <EOL> while error > tol and iter < iter_max : <EOL> assert error_grid . dtype == np . float64 <EOL> jacobi_relax_core [ griddim , blockdim , stream ] ( dA , dAnew , derror_grid ) <EOL> derror_grid . to_host ( stream ) <EOL> stream . synchronize ( ) <EOL> error = np . abs ( error_grid ) . max ( ) <EOL> tmp = dA <EOL> dA = dAnew <EOL> dAnew = tmp <EOL> if iter % <NUM_LIT:100> == <NUM_LIT:0> : <EOL> print ( "<STR_LIT>" % ( iter , error , time . time ( ) - timer ) ) <EOL> iter += <NUM_LIT:1> <EOL> runtime = time . time ( ) - timer <EOL> print ( "<STR_LIT>" % runtime ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> main ( ) </s>
