github.com/coyove/nj@v0.0.0-20221110084952-c7f8db1065c3/tests/bench/mandelbrot.go (about)

     1  package main
     2  
     3  import (
     4  	"bufio"
     5  	"flag"
     6  	"fmt"
     7  	"os"
     8  	"runtime"
     9  	"strconv"
    10  	"sync"
    11  	"sync/atomic"
    12  )
    13  
    14  const limit = 4.0      // abs(z) < 2
    15  const maxIter = 50     // number of iterations
    16  const defaultSize = 16 // bitmap size if not given as command-line argument
    17  
    18  var rows [][]byte
    19  var bytesPerRow int
    20  var initial_r []float64
    21  var initial_i []float64
    22  var zzz int
    23  
    24  // renderRow returns rendered row of pixels
    25  // pixels packed in one byte for PBM image
    26  func renderRow(y0 *int32) []byte {
    27  	var i, j, x int
    28  	var res, b byte
    29  	var Zr1, Zr2, Zi1, Zi2, Tr1, Tr2, Ti1, Ti2 float64
    30  
    31  	row := make([]byte, bytesPerRow)
    32  
    33  	for xByte := range row {
    34  		res = 0
    35  		Ci := initial_i[*y0]
    36  		for i = 0; i < 8; i += 2 {
    37  			x = xByte << 3 // * 8
    38  			Cr1 := initial_r[x+i]
    39  			Cr2 := initial_r[x+i+1]
    40  
    41  			Zr1 = Cr1
    42  			Zi1 = Ci
    43  
    44  			Zr2 = Cr2
    45  			Zi2 = Ci
    46  
    47  			b = 0
    48  
    49  			for j = 0; j < maxIter; j++ {
    50  				Tr1 = Zr1 * Zr1
    51  				Ti1 = Zi1 * Zi1
    52  				Zi1 = 2*Zr1*Zi1 + Ci
    53  				Zr1 = Tr1 - Ti1 + Cr1
    54  
    55  				if Tr1+Ti1 > limit {
    56  					b |= 2
    57  					if b == 3 {
    58  						break
    59  					}
    60  				}
    61  
    62  				Tr2 = Zr2 * Zr2
    63  				Ti2 = Zi2 * Zi2
    64  				Zi2 = 2*Zr2*Zi2 + Ci
    65  				Zr2 = Tr2 - Ti2 + Cr2
    66  
    67  				if Tr2+Ti2 > limit {
    68  					b |= 1
    69  					if b == 3 {
    70  						break
    71  					}
    72  				}
    73  			}
    74  			res = (res << 2) | b
    75  		}
    76  		row[xByte] = ^res
    77  	}
    78  	fmt.Printf("%02x\n", row)
    79  	return row
    80  }
    81  
    82  var yAt int32 = -1
    83  
    84  func renderRows(wg *sync.WaitGroup, s32 int32) {
    85  	var y int32
    86  	for y = atomic.AddInt32(&yAt, 1); y < s32; y = atomic.AddInt32(&yAt, 1) {
    87  		rows[y] = renderRow(&y)
    88  	}
    89  	wg.Done()
    90  }
    91  
    92  func main() {
    93  	pool := runtime.NumCPU() * 2
    94  	pool = 1
    95  	runtime.GOMAXPROCS(pool)
    96  
    97  	// hashGet input, if any...
    98  	size := defaultSize
    99  	flag.Parse()
   100  	if flag.NArg() > 0 {
   101  		size, _ = strconv.Atoi(flag.Arg(0))
   102  	}
   103  
   104  	bytesPerRow = size >> 3
   105  
   106  	// Precompute the initial real and imaginary values for each x and y
   107  	// coordinate in the image.
   108  	initial_r = make([]float64, size)
   109  	initial_i = make([]float64, size)
   110  	inv := 2.0 / float64(size)
   111  	for xy := 0; xy < size; xy++ {
   112  		i := inv * float64(xy)
   113  		initial_r[xy] = i - 1.5
   114  		initial_i[xy] = i - 1.0
   115  	}
   116  
   117  	rows = make([][]byte, size)
   118  
   119  	/* Wait group for finish */
   120  	wg := new(sync.WaitGroup)
   121  	wg.Add(pool)
   122  
   123  	// start pool workers, and assign all work
   124  	for i := 0; i < pool; i++ {
   125  		go renderRows(wg, int32(size))
   126  	}
   127  
   128  	/* wait for the file workers to finish, then write */
   129  	wg.Wait()
   130  
   131  	out := bufio.NewWriter(os.Stdout)
   132  	defer out.Flush()
   133  	fmt.Fprintf(out, "P4\n%d %d\n", size, size)
   134  
   135  	for y := 0; y < size; y++ {
   136  		out.Write(rows[y])
   137  	}
   138  }