Merge pull request #23 from rust-math/package

Package subcommand for intel-mkl-tool

Author: termoshtt

ci/Dockerfile

@@ -1,15 +1,8 @@
 FROM rust:latest
 LABEL maintainer "Toshiki Teramura <[email protected]>"
 
-RUN apt-get update \
- && apt-get install -y wget apt-transport-https \
- && wget https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS-2019.PUB \
- && apt-key add GPG-PUB-KEY-INTEL-SW-PRODUCTS-2019.PUB \
- && wget https://apt.repos.intel.com/setup/intelproducts.list -O /etc/apt/sources.list.d/intelproducts.list \
- && apt-get update \
- && apt-get install -y intel-mkl-core-rt-2019.3 \
- && apt-get clean \
- && rm -rf /var/lib/apt/lists/*
+COPY install.sh /
+RUN /install.sh && apt-get clean && rm -rf /var/lib/apt/lists/*
 
 ENV PKG_CONFIG_PATH /opt/intel/compilers_and_libraries/linux/mkl/bin/pkgconfig
 ENV LD_LIBRARY_PATH /opt/intel/compilers_and_libraries/linux/mkl/lib/intel64

ci/install.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+# Intel MKL installer for ubuntu/debian
+set -eux
+
+apt-get update
+apt-get install -y wget apt-transport-https
+curl -sfL https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS-2019.PUB | apt-key add -
+curl -sfL https://apt.repos.intel.com/setup/intelproducts.list -o /etc/apt/sources.list.d/intelproducts.list
+apt-get update
+apt-get install -y intel-mkl-core-rt-2019.5  # Please update version number

intel-mkl-tool/Cargo.toml

@@ -13,12 +13,13 @@ curl = "0.4.25"
 failure = "0.1.6"
 pkg-config = "0.3.17"
 tar = "0.4.26"
-xz2 = "0.1.6"
 log = "0.4.8"
-dirs = "*"
+dirs = "2.0.2"
+glob = "0.3.0"
+zstd = "0.5.1"
 
 # CLI
-structopt  = { version = "0.3.4", optional = true }
+structopt  = { version = "0.3.5", optional = true }
 env_logger = { version = "0.7.1", optional = true }
 
 [[bin]]

intel-mkl-tool/src/cli.rs

@@ -10,12 +10,16 @@ enum Opt {
         /// Install destination. Default is `$XDG_DATA_HOME/intel-mkl-tool`
         path: Option<PathBuf>,
     },
+
     /// Seek Intel-MKL library
     ///
     /// 1. pkg-config
     /// 2. `$XDG_DATA_HOME/intel-mkl-tool`
     /// will be sought.
     Seek {},
+
+    /// Package Intel MKL libraries into an archive
+    Package { path: PathBuf },
 }
 
 fn main() -> Fallible<()> {
@@ -33,6 +37,7 @@ fn main() -> Fallible<()> {
             };
             intel_mkl_tool::download(&path)?;
         }
+
         Opt::Seek {} => {
             if let Some(path) = intel_mkl_tool::seek_pkg_config() {
                 println!("{}", path.display());
@@ -44,6 +49,10 @@ fn main() -> Fallible<()> {
             }
             bail!("Intel-MKL not found.");
         }
+
+        Opt::Package { path } => {
+            let _out = intel_mkl_tool::package(&path)?;
+        }
     }
     Ok(())
 }

intel-mkl-tool/src/lib.rs

@@ -1,33 +1,49 @@
 use curl::easy::Easy;
 use failure::*;
+use glob::glob;
 use log::*;
 use std::{
     fs,
-    io::{self, Write},
+    io::{self, BufRead, Write},
     path::*,
 };
 
 const S3_ADDR: &'static str = "https://s3-ap-northeast-1.amazonaws.com/rust-intel-mkl";
 
-#[cfg(target_os = "linux")]
+#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
 mod mkl {
-    pub const ARCHIVE: &'static str = "mkl_linux.tar.xz";
+    pub const ARCHIVE: &'static str = "mkl_linux64";
     pub const EXT: &'static str = "so";
     pub const PREFIX: &'static str = "lib";
+    pub const VERSION_YEAR: u32 = 2019;
+    pub const VERSION_UPDATE: u32 = 5;
 }
 
-#[cfg(target_os = "macos")]
+#[cfg(all(target_os = "macos", target_arch = "x86_64"))]
 mod mkl {
-    pub const ARCHIVE: &'static str = "mkl_osx.tar.xz";
+    pub const ARCHIVE: &'static str = "mkl_macos64";
     pub const EXT: &'static str = "dylib";
     pub const PREFIX: &'static str = "lib";
+    pub const VERSION_YEAR: u32 = 2019;
+    pub const VERSION_UPDATE: u32 = 3;
 }
 
-#[cfg(target_os = "windows")]
+#[cfg(all(target_os = "windows", target_arch = "x86_64"))]
 mod mkl {
-    pub const ARCHIVE: &'static str = "mkl_windows64.tar.xz";
+    pub const ARCHIVE: &'static str = "mkl_windows64";
     pub const EXT: &'static str = "lib";
     pub const PREFIX: &'static str = "";
+    pub const VERSION_YEAR: u32 = 2019;
+    pub const VERSION_UPDATE: u32 = 5;
+}
+
+pub fn archive_filename() -> String {
+    format!(
+        "{}_{}_{}.tar.zst",
+        mkl::ARCHIVE,
+        mkl::VERSION_YEAR,
+        mkl::VERSION_UPDATE
+    )
 }
 
 pub fn home_library_path() -> PathBuf {
@@ -61,13 +77,14 @@ pub fn download(out_dir: &Path) -> Fallible<()> {
         bail!("Not a directory: {}", out_dir.display());
     }
 
-    let archive = out_dir.join(mkl::ARCHIVE);
+    let archive = out_dir.join(archive_filename());
     if !archive.exists() {
-        info!("Download archive from AWS S3: {}/{}", S3_ADDR, mkl::ARCHIVE);
+        let url = format!("{}/{}", S3_ADDR, archive_filename());
+        info!("Download archive from AWS S3: {}", url);
         let f = fs::File::create(&archive)?;
         let mut buf = io::BufWriter::new(f);
         let mut easy = Easy::new();
-        easy.url(&format!("{}/{}", S3_ADDR, mkl::ARCHIVE))?;
+        easy.url(&url)?;
         easy.write_function(move |data| Ok(buf.write(data).unwrap()))?;
         easy.perform()?;
         assert!(archive.exists());
@@ -78,7 +95,7 @@ pub fn download(out_dir: &Path) -> Fallible<()> {
     let core = out_dir.join(format!("{}mkl_core.{}", mkl::PREFIX, mkl::EXT));
     if !core.exists() {
         let f = fs::File::open(&archive)?;
-        let de = xz2::read::XzDecoder::new(f);
+        let de = zstd::stream::read::Decoder::new(f)?;
         let mut arc = tar::Archive::new(de);
         arc.unpack(&out_dir)?;
         assert!(core.exists());
@@ -87,3 +104,63 @@ pub fn download(out_dir: &Path) -> Fallible<()> {
     }
     Ok(())
 }
+
+// Read mkl_version.h to get MKL version (e.g. 2019.5)
+fn get_mkl_version(version_header: &Path) -> Fallible<(u32, u32)> {
+    if !version_header.exists() {
+        bail!("MKL Version file not found: {}", version_header.display());
+    }
+    let f = fs::File::open(version_header)?;
+    let f = io::BufReader::new(f);
+    let mut year = 0;
+    let mut update = 0;
+    for line in f.lines() {
+        if let Ok(line) = line {
+            if !line.starts_with("#define") {
+                continue;
+            }
+            let ss: Vec<&str> = line.split(" ").collect();
+            match ss[1] {
+                "__INTEL_MKL__" => year = ss[2].parse()?,
+                "__INTEL_MKL_UPDATE__" => update = ss[2].parse()?,
+                _ => continue,
+            }
+        }
+    }
+    if year == 0 || update == 0 {
+        bail!("Cannot determine MKL versions");
+    }
+    Ok((year, update))
+}
+
+pub fn package(mkl_path: &Path) -> Fallible<PathBuf> {
+    if !mkl_path.exists() {
+        bail!("MKL directory not found: {}", mkl_path.display());
+    }
+    let (year, update) = get_mkl_version(&mkl_path.join("include/mkl_version.h"))?;
+    info!("Intel MKL version: {}.{}", year, update);
+    let out = PathBuf::from(archive_filename());
+    info!("Create archive: {}", out.display());
+
+    let shared_libs: Vec<_> = glob(
+        mkl_path
+            .join(format!("lib/intel64/*.{}", mkl::EXT))
+            .to_str()
+            .unwrap(),
+    )?
+    .map(|path| path.unwrap())
+    .collect();
+    let f = fs::File::create(&out)?;
+    let buf = io::BufWriter::new(f);
+    let zstd = zstd::stream::write::Encoder::new(buf, 6)?;
+    let mut ar = tar::Builder::new(zstd);
+    ar.mode(tar::HeaderMode::Deterministic);
+    for lib in &shared_libs {
+        info!("Add {}", lib.display());
+        ar.append_path_with_name(lib, lib.file_name().unwrap())?;
+    }
+    let zstd = ar.into_inner()?;
+    zstd.finish()?;
+
+    Ok(out)
+}