Kirill Shakirov
400 lines
11 KiB
Rust
400 lines
11 KiB
Rust
extern crate ocl;
|
|
use ocl::{
|
|
Buffer, Context, Device, DeviceType, Kernel, Platform, Program, Queue, SpatialDims, flags,
|
|
};
|
|
use serde::de::value::Error;
|
|
use std::{
|
|
io,
|
|
str::{self, FromStr},
|
|
};
|
|
|
|
use crate::client_config::{AppConfig, DevConf};
|
|
|
|
mod client_config;
|
|
mod num_utils;
|
|
|
|
|
|
fn dev_type_from_str(s: &str) -> Result<flags::DeviceType, ()> {
|
|
match s {
|
|
"CPU" => Ok(flags::DeviceType::CPU),
|
|
"GPU" => Ok(flags::DeviceType::GPU),
|
|
"ALL" => Ok(flags::DeviceType::ALL),
|
|
"CUSTOM" => Ok(flags::DeviceType::CUSTOM),
|
|
"ACCELERATOR" => Ok(flags::DeviceType::ACCELERATOR),
|
|
"DEFAULT" => Ok(flags::DeviceType::DEFAULT),
|
|
_ => Err(()),
|
|
}
|
|
}
|
|
|
|
fn str_or_empty(r: ocl::error::Result<String>) -> String {
|
|
match r {
|
|
Ok(s) => s,
|
|
Err(_) => "".to_string(),
|
|
}
|
|
}
|
|
|
|
fn print_devices(dev_list: &Vec<(Device, Platform)>) {
|
|
let mut i = 0;
|
|
for (dev, plt) in dev_list.iter() {
|
|
let dev_name = str_or_empty(dev.name());
|
|
let plt_name = str_or_empty(plt.name());
|
|
println!("({i}) device: \"{dev_name}\" ----- platorm: \"{plt_name}\"");
|
|
i += 1;
|
|
}
|
|
}
|
|
|
|
fn choose_devices(devices_num: usize) -> Result<Vec<usize>, String> {
|
|
println!("Please input desired devices to use as a white space separated list of numbers.");
|
|
let mut result: Vec<usize> = Vec::new();
|
|
|
|
let mut s_devs_nums = String::new();
|
|
|
|
io::stdin()
|
|
.read_line(&mut s_devs_nums)
|
|
.expect("Failed to read line");
|
|
|
|
for s_dev_num in s_devs_nums.split(' ') {
|
|
let dev_num: usize = match s_dev_num.trim().parse() {
|
|
Ok(num) => num,
|
|
Err(_) => return Err("You must input a number from device list.".to_string()),
|
|
};
|
|
if dev_num >= devices_num {
|
|
return Err("You must input a number from device list.".to_string());
|
|
};
|
|
result.push(dev_num);
|
|
}
|
|
return Ok(result);
|
|
}
|
|
|
|
fn list_devices(dev_type: DeviceType) -> Vec<(Device, Platform)> {
|
|
let platforms = Platform::list();
|
|
let mut devices: Vec<(Device, Platform)> = Vec::new();
|
|
for plt in platforms.iter() {
|
|
//let plat_name = str_or_empty(plt.name());
|
|
let list_res = Device::list(plt, Some(dev_type));
|
|
match list_res {
|
|
Ok(dev_l) => devices.extend(dev_l.iter().map(|dev| (*dev, plt.clone()))),
|
|
Err(_) => {}
|
|
}
|
|
}
|
|
return devices;
|
|
}
|
|
|
|
// fn init_devices(
|
|
// devices: Vec<(Device, Platform, DevConfig)>,
|
|
// kern_name: String,
|
|
// prog_src: String,
|
|
// inc_dirs: Vec<String>,
|
|
// ) -> Vec<ExecContext> {
|
|
// let mut contexts: Vec<ExecContext> = Vec::with_capacity(devices.len());
|
|
// for (dev, plt, dev_cfg) in devices {
|
|
// let ctx = match Context::builder()
|
|
// .platform(plt)
|
|
// .devices(dev.clone())
|
|
// .build()
|
|
// {
|
|
// Ok(c) => c,
|
|
// Err(_) => continue,
|
|
// };
|
|
|
|
// let prg = match Program::builder().devices(dev).src(&prog_src).build(&ctx) {
|
|
// Ok(p) => p,
|
|
// Err(_) => continue,
|
|
// };
|
|
|
|
// let queue = match Queue::new(&ctx, dev, None) {
|
|
// Ok(q) => q,
|
|
// Err(_) => continue,
|
|
// };
|
|
|
|
// // Create Buffers:
|
|
// let start_key_b = match Buffer::<u32>::builder()
|
|
// .queue(queue.clone())
|
|
// .flags(flags::MEM_READ_ONLY)
|
|
// .len(8)
|
|
// .fill_val(0u32)
|
|
// .build()
|
|
// {
|
|
// Ok(buf) => buf,
|
|
// Err(_) => continue,
|
|
// };
|
|
|
|
// let u_data_b = match Buffer::<u32>::builder()
|
|
// .queue(queue.clone())
|
|
// .flags(flags::MEM_READ_ONLY)
|
|
// .len(4)
|
|
// .fill_val(0u32)
|
|
// .build()
|
|
// {
|
|
// Ok(buf) => buf,
|
|
// Err(_) => continue,
|
|
// };
|
|
|
|
// let enc_data_b = match Buffer::<u32>::builder()
|
|
// .queue(queue.clone())
|
|
// .flags(flags::MEM_READ_ONLY)
|
|
// .len(4)
|
|
// .fill_val(0u32)
|
|
// .build()
|
|
// {
|
|
// Ok(buf) => buf,
|
|
// Err(_) => continue,
|
|
// };
|
|
|
|
// let key_found_b = match Buffer::<u32>::builder()
|
|
// .queue(queue.clone())
|
|
// .flags(flags::MEM_WRITE_ONLY)
|
|
// .len(1)
|
|
// .fill_val(0u32)
|
|
// .build()
|
|
// {
|
|
// Ok(buf) => buf,
|
|
// Err(_) => continue,
|
|
// };
|
|
|
|
// // (3) Create a kernel with arguments matching those in the source above:
|
|
// let kernel = match Kernel::builder()
|
|
// .program(&prg)
|
|
// .name(&kern_name)
|
|
// .queue(queue.clone())
|
|
// .global_work_size(dev_cfg.global_work_size)
|
|
// .arg(&start_key_b)
|
|
// .arg(&u_data_b)
|
|
// .arg(&enc_data_b)
|
|
// .arg(&key_found_b)
|
|
// .build()
|
|
// {
|
|
// Ok(kern) => kern,
|
|
// Err(_) => continue,
|
|
// };
|
|
|
|
// contexts.push(ExecContext {
|
|
// cfg: dev_cfg,
|
|
// ctx: ctx,
|
|
// kernel: kernel,
|
|
// prog: prg,
|
|
// queue: queue,
|
|
// buffers: CtxBuffers {
|
|
// start_key: start_key_b,
|
|
// u_data: u_data_b,
|
|
// enc_data: enc_data_b,
|
|
// key_found: key_found_b,
|
|
// },
|
|
// });
|
|
// }
|
|
// return contexts;
|
|
// }
|
|
|
|
fn dev_sel_dialog(all_devices: &Vec<(Device, Platform)>) -> Vec<usize> {
|
|
let devs_nums = loop {
|
|
print_devices(&all_devices);
|
|
match choose_devices(all_devices.len()) {
|
|
Ok(value) => break value,
|
|
Err(exc) => {
|
|
println!("Error! {exc}\n")
|
|
}
|
|
}
|
|
};
|
|
return devs_nums;
|
|
}
|
|
|
|
fn get_devices_conf(file_name: &str) -> Result<(Vec<(Device, Platform)>, AppConfig), String> {
|
|
let dev_type = dev_type_from_str("GPU").expect("Unexpected device type!");
|
|
|
|
// Get devices to be used for key search
|
|
let all_devices: Vec<(Device, Platform)> = list_devices(dev_type);
|
|
if all_devices.len() == 0 {
|
|
return Err("Cannot find any usable devices.".to_string());
|
|
};
|
|
|
|
let app_conf = match client_config::load_config(file_name) {
|
|
Ok(readed_config) => {
|
|
let dev_found = all_devices
|
|
.iter()
|
|
.filter(|dp| readed_config.device_exist(&dp.0))
|
|
.count();
|
|
if dev_found < readed_config.devices.len() {
|
|
println!("Devices from config not found in the system!");
|
|
let devs_nums = dev_sel_dialog(&all_devices);
|
|
let res = AppConfig::from_dev_list(&all_devices, devs_nums);
|
|
client_config::save_config(file_name, &res);
|
|
res
|
|
} else {
|
|
readed_config
|
|
}
|
|
}
|
|
Err(_) => {
|
|
println!("Cannot find config file {}", file_name);
|
|
let devs_nums = dev_sel_dialog(&all_devices);
|
|
AppConfig::from_dev_list(&all_devices, devs_nums)
|
|
}
|
|
};
|
|
|
|
let selected_devs = all_devices
|
|
.iter()
|
|
.filter(|dp| app_conf.device_exist(&dp.0))
|
|
.cloned()
|
|
.collect();
|
|
|
|
return Ok((selected_devs, app_conf));
|
|
}
|
|
|
|
struct CtxBuffers {
|
|
batch_size: u32,
|
|
tweak_i: u64,
|
|
tweak_j: u32,
|
|
start_key: Buffer<u32>,
|
|
uenc_data: Buffer<u32>,
|
|
target_data: Buffer<u32>,
|
|
key_found: Buffer<u32>,
|
|
}
|
|
|
|
struct ExecData {
|
|
start_key: Vec<u32>,
|
|
uenc_data: Vec<u32>,
|
|
target_data: Vec<u32>,
|
|
key_found: Vec<u32>,
|
|
batch_size: u32,
|
|
work_size: usize,
|
|
}
|
|
struct ExecContext {
|
|
ctx: Context,
|
|
kernel: Kernel,
|
|
prog: Program,
|
|
queue: Queue,
|
|
buffers: CtxBuffers,
|
|
exec_data: ExecData,
|
|
}
|
|
|
|
fn init_program(
|
|
cl_device: Device,
|
|
cl_platform: Platform,
|
|
cl_src: &str,
|
|
cl_cmplr_opt: &str,
|
|
) -> Result<(Context, Program, Queue), ocl::Error> {
|
|
let cl_context = Context::builder()
|
|
.platform(cl_platform)
|
|
.devices(cl_device.clone())
|
|
.build()?;
|
|
|
|
let cl_program = Program::builder()
|
|
.devices(cl_device)
|
|
.src(cl_src)
|
|
.cmplr_opt(cl_cmplr_opt)
|
|
.build(&cl_context)
|
|
.unwrap();
|
|
|
|
let cl_queue: Queue = Queue::new(&cl_context, cl_device, None)?;
|
|
|
|
return Ok((cl_context, cl_program, cl_queue));
|
|
}
|
|
|
|
fn init_buffers(cl_queue: Queue) -> Result<CtxBuffers, ocl::Error> {
|
|
let cl_buffer_start_key = Buffer::<u32>::builder()
|
|
.queue(cl_queue.clone())
|
|
.flags(flags::MEM_READ_ONLY)
|
|
.len(8)
|
|
.fill_val(0u32)
|
|
.build()?;
|
|
|
|
let cl_buffer_uenc_data = Buffer::<u32>::builder()
|
|
.queue(cl_queue.clone())
|
|
.flags(flags::MEM_READ_ONLY)
|
|
.len(4)
|
|
.fill_val(0u32)
|
|
.build()?;
|
|
|
|
let cl_buffer_target_data = Buffer::<u32>::builder()
|
|
.queue(cl_queue.clone())
|
|
.flags(flags::MEM_READ_ONLY)
|
|
.len(4)
|
|
.fill_val(0u32)
|
|
.build()?;
|
|
|
|
let cl_buffer_key_found = Buffer::<u32>::builder()
|
|
.queue(cl_queue.clone())
|
|
.flags(flags::MEM_WRITE_ONLY)
|
|
.len(9)
|
|
.fill_val(0u32)
|
|
.build()?;
|
|
|
|
Ok(CtxBuffers {
|
|
batch_size: 0,
|
|
tweak_i: 0,
|
|
tweak_j: 0,
|
|
start_key: cl_buffer_start_key,
|
|
uenc_data: cl_buffer_uenc_data,
|
|
target_data: cl_buffer_target_data,
|
|
key_found: cl_buffer_key_found,
|
|
})
|
|
}
|
|
|
|
fn init_kernel(
|
|
work_size: usize,
|
|
cl_program: Program,
|
|
cl_queue: Queue,
|
|
buffs: &CtxBuffers,
|
|
) -> Result<Kernel, ocl::Error> {
|
|
Kernel::builder()
|
|
.program(&cl_program)
|
|
.name("search_key")
|
|
.queue(cl_queue.clone())
|
|
.global_work_size(work_size)
|
|
.arg(&buffs.batch_size)
|
|
.arg(&buffs.tweak_i)
|
|
.arg(&buffs.tweak_j)
|
|
.arg(&buffs.start_key)
|
|
.arg(&buffs.uenc_data)
|
|
.arg(&buffs.target_data)
|
|
.arg(&buffs.key_found)
|
|
.build()
|
|
}
|
|
|
|
fn do_work(ex_ctx: &mut ExecContext) -> Result<bool, ocl::Error> {
|
|
ex_ctx.buffers.batch_size = ex_ctx.exec_data.batch_size;
|
|
ex_ctx
|
|
.buffers
|
|
.start_key
|
|
.cmd()
|
|
.queue(&ex_ctx.queue)
|
|
.offset(0)
|
|
.write(&ex_ctx.exec_data.start_key)
|
|
.enq()?;
|
|
|
|
// (4) Run the kernel
|
|
unsafe {
|
|
ex_ctx
|
|
.kernel
|
|
.cmd()
|
|
.queue(&ex_ctx.queue)
|
|
.global_work_size(ex_ctx.exec_data.work_size)
|
|
.enq()?;
|
|
}
|
|
|
|
ex_ctx
|
|
.buffers
|
|
.key_found
|
|
.cmd()
|
|
.queue(&ex_ctx.queue)
|
|
.offset(0)
|
|
.read(&mut ex_ctx.exec_data.key_found)
|
|
.enq()?;
|
|
|
|
if ex_ctx.exec_data.key_found[0] == 0{Ok(false)}
|
|
else {Ok(true)}
|
|
}
|
|
|
|
fn main() {
|
|
println!("Hello, world nya!");
|
|
//use ocl::{Buffer, Context, Device, Kernel, Platform, Program, Queue, flags};
|
|
let devices = get_devices_conf("test.json");
|
|
println!("{:?}", devices);
|
|
|
|
// let devices: Vec<_> = platforms.iter().flat_map(|p| Device::list(p, Some(dev_type)).iter()).collect();
|
|
// let device = Device::first(platform)?;
|
|
// let context = Context::builder()
|
|
// .platform(platform)
|
|
// .devices(device.clone())
|
|
// .build()?;
|
|
}
|