add parallel binary edge list shuffler script (#67)

scripts/CMakeLists.txt

@@ -7,6 +7,3 @@ target_link_libraries(divide-into-batches PRIVATE Boost::program_options)
 
 add_executable(count-batched count_batched.cpp)
 target_link_libraries(count-batched PRIVATE Boost::program_options)
-
-add_executable(shufbel shufbel.cpp)
-target_link_libraries(shufbel PRIVATE Boost::program_options)

scripts/shufbel/go.mod

@@ -0,0 +1,5 @@
+module shufbel
+
+go 1.18
+
+require github.com/vmunoz82/shuffle v1.0.2 // indirect

scripts/shufbel/go.sum

@@ -0,0 +1,2 @@
+github.com/vmunoz82/shuffle v1.0.2 h1:9QxvJN9F1OdsnRVrPu1+N1F+IU14PPtHATtvlxTSDVI=
+github.com/vmunoz82/shuffle v1.0.2/go.mod h1:MvogPEkxTyK+LU0v+nYDCzWiD8Jlu/rocUCJOvs8mcA=

scripts/shufbel/shufbel.go

@@ -0,0 +1,128 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"math/rand"
+	"os"
+	"runtime"
+	"sync"
+	"syscall"
+	"unsafe"
+
+	"github.com/vmunoz82/shuffle"
+)
+
+/* Function for each round, could be anything don't need to be reversible */
+func roundFunction(v, key shuffle.FeistelWord) shuffle.FeistelWord {
+	return (v * 941083987) ^ (key >> (v & 7) * 104729)
+}
+
+// main function
+func main() {
+	// parse CLI options
+	threads := flag.Int("threads", runtime.NumCPU(), "number of threads")
+	rseed := flag.Int64("rseed", 0, "random seed")
+	flag.Parse()
+
+	// limit go runtime threads
+	runtime.GOMAXPROCS(*threads)
+	pageSize := syscall.Getpagesize()
+
+	files := flag.Args()
+	if len(files) != 2 {
+		fmt.Println("Usage: shufbel [options] input output ...")
+		flag.PrintDefaults()
+		return
+	}
+	input_file := files[0]
+	output_file := files[1]
+
+	// mmap input file
+	fin, err := os.Open(input_file)
+	if err != nil {
+		fmt.Printf("Error opening file %s: %v\n", input_file, err)
+		return
+	}
+	defer fin.Close()
+
+	fout, err := os.OpenFile(output_file, os.O_RDWR|os.O_CREATE, 0644)
+	if err != nil {
+		fmt.Printf("Error opening file %s: %v\n", output_file, err)
+		return
+	}
+
+	fi, err := fin.Stat()
+	if err != nil {
+		fmt.Printf("Error getting file info: %v\n", err)
+	}
+	sizeBytes := uint64(fi.Size())
+	if sizeBytes == 0 {
+		fmt.Printf("Error: file size is 0: %s\n", input_file)
+		return
+	}
+	if sizeBytes%16 != 0 {
+		fmt.Printf("Error: file size is not divisible by 16")
+		return
+	}
+
+	fout.Truncate(int64(sizeBytes))
+
+	data, err := syscall.Mmap(int(fin.Fd()), 0, int(sizeBytes), syscall.PROT_READ, syscall.MAP_PRIVATE|syscall.MAP_NORESERVE)
+	if err != nil {
+		fmt.Printf("Error mmaping file %s: %v\n", input_file, err)
+		return
+	}
+	if uintptr(unsafe.Pointer(&data[0]))%uintptr(pageSize) != 0 {
+		fmt.Println("mmap returned a non-page-aligned address: %p", uintptr(unsafe.Pointer(&data[0])))
+		return
+	}
+	// madv willneed
+	err = syscall.Madvise(data, syscall.MADV_WILLNEED)
+	if err != nil {
+		fmt.Printf("Error madvise: %v\n", err)
+		return
+	}
+
+	// mmap output file
+	output, err := syscall.Mmap(int(fout.Fd()), 0, int(sizeBytes), syscall.PROT_WRITE, syscall.MAP_SHARED)
+	if err != nil {
+		fmt.Printf("Error mmaping file %s: %v\n", output_file, err)
+		return
+	}
+
+	num_edges := sizeBytes / 16
+
+	rand.Seed(*rseed)
+	// keys should be an array of 4 random uint64s
+	keys := []shuffle.FeistelWord{shuffle.FeistelWord(rand.Uint64()), shuffle.FeistelWord(rand.Uint64()), shuffle.FeistelWord(rand.Uint64()), shuffle.FeistelWord(rand.Uint64())}
+
+	var wg sync.WaitGroup
+	// spawn a goroutine for every page, since each page can fault
+	// todo (meyer): should we limit the number of concurrent goroutines?
+	for i := uint64(0); i < num_edges; i += uint64(pageSize / 16) {
+		wg.Add(1)
+		go func(i uint64) {
+			defer wg.Done()
+			fn := shuffle.NewFeistel(keys, roundFunction)
+			for j := i; j < i+uint64(pageSize/16) && j < num_edges; j++ {
+				src, _ := shuffle.GetIndex(shuffle.FeistelWord(j), shuffle.FeistelWord(num_edges), fn)
+				// move 16 bytes from data to output
+				copy(output[j*16:], data[src*16:src*16+16])
+			}
+		}(i)
+	}
+	wg.Wait()
+
+	err = syscall.Munmap(data)
+	if err != nil {
+		fmt.Printf("Error unmapping file %s: %v\n", input_file, err)
+		return
+	}
+
+	err = syscall.Munmap(output)
+	if err != nil {
+		fmt.Printf("Error unmapping file %s: %v\n", output_file, err)
+		return
+	}
+}