compdist - project report
main goal
- explore the intersection betweenc compilers and distributed computing, trying to enrich compilers with domain specific information about the target system, e.g., topology, nodes capacity, etc.
progress
- the basic idea is to have a pass to optimize affine -> mpi, starting with a basic partition (A+B) in two ranks, processing a half of the data on each process []
tasks
- clean mlir-tutorial code []
- create a basic pass using cpp mechanism []
- update basic version with more ranks (e.g., -n-ranks=10) []
- create a new version where the number of ranks is passed as an option and the pass distributes the data uniformly to every rank []
Examples
A + B (partition in halfs)
Input:
func.func @add_arrays(%A: memref<100xf32>, %B: memref<100xf32>, %C: memref<100xf32>) {
affine.for %i = 0 to 100 {
%a = memref.load %A[%i] : memref<100xf32>
%b = memref.load %B[%i] : memref<100xf32>
%sum = arith.addf %a, %b : f32
memref.store %sum, %C[%i] : memref<100xf32>
}
return
}
affine.for %i = 0 to 100 {
%a = memref.load %A[%i] : memref<100xf32>
%b = memref.load %B[%i] : memref<100xf32>
%sum = arith.addf %a, %b : f32
memref.store %sum, %C[%i] : memref<100xf32>
}
return
}
Output:
func.func @add_arrays_mpi(%A: memref<100xf32>, %B: memref<100xf32>, %C: memref<100xf32>) {
%c0 = arith.constant 0 : i32
%c1 = arith.constant 1 : i32
%c50 = arith.constant 50 : index // Half of 100
// Get MPI rank
%init_err = mpi.init : !mpi.retval
%rank_err, %rank = mpi.comm_rank : !mpi.retval, i32
// Process based on rank
%zero = arith.constant 0 : i32
%is_rank_zero = arith.cmpi eq, %rank, %zero : i32
scf.if %is_rank_zero {
// Process 0 sends first half of data
%send_err1 = mpi.send(%A, %c1, %c0) : memref<100xf32>, i32, i32 -> !mpi.retval
%send_err2 = mpi.send(%B, %c1, %c0) : memref<100xf32>, i32, i32 -> !mpi.retval
// Process local half
affine.for %i = 50 to 100 {
%a = memref.load %A[%i] : memref<100xf32>
%b = memref.load %B[%i] : memref<100xf32>
%sum = arith.addf %a, %b : f32
memref.store %sum, %C[%i] : memref<100xf32>
}
// Receive processed first half
%recv_err = mpi.recv(%C, %c1, %c0) : memref<100xf32>, i32, i32 -> !mpi.retval
} else {
// Process 1 receives and processes first half
%local_a = memref.alloc() : memref<100xf32>
%local_b = memref.alloc() : memref<100xf32>
%local_result = memref.alloc() : memref<100xf32>
%recv_err1 = mpi.recv(%local_a, %c0, %c0) : memref<100xf32>, i32, i32 -> !mpi.retval
%recv_err2 = mpi.recv(%local_b, %c0, %c0) : memref<100xf32>, i32, i32 -> !mpi.retval
affine.for %i = 0 to 50 {
%a = memref.load %local_a[%i] : memref<100xf32>
%b = memref.load %local_b[%i] : memref<100xf32>
%sum = arith.addf %a, %b : f32
memref.store %sum, %local_result[%i] : memref<100xf32>
}
%send_err = mpi.send(%local_result, %c0, %c0) : memref<100xf32>, i32, i32 -> !mpi.retval
memref.dealloc %local_a : memref<100xf32>
memref.dealloc %local_b : memref<100xf32>
memref.dealloc %local_result : memref<100xf32>
}
return
}
%c0 = arith.constant 0 : i32
%c1 = arith.constant 1 : i32
%c50 = arith.constant 50 : index // Half of 100
// Get MPI rank
%init_err = mpi.init : !mpi.retval
%rank_err, %rank = mpi.comm_rank : !mpi.retval, i32
// Process based on rank
%zero = arith.constant 0 : i32
%is_rank_zero = arith.cmpi eq, %rank, %zero : i32
scf.if %is_rank_zero {
// Process 0 sends first half of data
%send_err1 = mpi.send(%A, %c1, %c0) : memref<100xf32>, i32, i32 -> !mpi.retval
%send_err2 = mpi.send(%B, %c1, %c0) : memref<100xf32>, i32, i32 -> !mpi.retval
// Process local half
affine.for %i = 50 to 100 {
%a = memref.load %A[%i] : memref<100xf32>
%b = memref.load %B[%i] : memref<100xf32>
%sum = arith.addf %a, %b : f32
memref.store %sum, %C[%i] : memref<100xf32>
}
// Receive processed first half
%recv_err = mpi.recv(%C, %c1, %c0) : memref<100xf32>, i32, i32 -> !mpi.retval
} else {
// Process 1 receives and processes first half
%local_a = memref.alloc() : memref<100xf32>
%local_b = memref.alloc() : memref<100xf32>
%local_result = memref.alloc() : memref<100xf32>
%recv_err1 = mpi.recv(%local_a, %c0, %c0) : memref<100xf32>, i32, i32 -> !mpi.retval
%recv_err2 = mpi.recv(%local_b, %c0, %c0) : memref<100xf32>, i32, i32 -> !mpi.retval
affine.for %i = 0 to 50 {
%a = memref.load %local_a[%i] : memref<100xf32>
%b = memref.load %local_b[%i] : memref<100xf32>
%sum = arith.addf %a, %b : f32
memref.store %sum, %local_result[%i] : memref<100xf32>
}
%send_err = mpi.send(%local_result, %c0, %c0) : memref<100xf32>, i32, i32 -> !mpi.retval
memref.dealloc %local_a : memref<100xf32>
memref.dealloc %local_b : memref<100xf32>
memref.dealloc %local_result : memref<100xf32>
}
return
}