Memory-mapped Utility for PECOS XLinear Model

[weiliw-amz]

Issue #, if available:
N/A

Description of changes:

This pull request will consist of following 3 commits:

• Refactor chunked matrix for PECOS XLinear model inference
  • Concatenated original chunked matrix's fragmented memory allocation
    • For accommodating subsequent memory-mapped utility module.
    • This change increases time cost of making chunked matrix by 10%~15% for large models (>50G), but is necessary and cannot be avoided.
  • Reduced memory footprint of making chunked matrix
• Memory-mapped utility module
  • Easy to use, well-encapsulated tool designed for dumping/loading arbitrary PECOS model
• Memory-mapped PECOS XLinear model
  • Greatly reduce loading time.
  • Ideal for large models that user want to quickly try a few inferences without waiting for loading full model into memory.
  • Also capable for large model inference that could not be stored in memory.

Usage:
User needs to have a XLinear Model saved on disk (in original .npz format), and manually compile into mmap format by calling compile_mmap_model:

import sys
from pecos.xmc.xlinear.model import XLinearModel


npz_model_path = f"/path/to/xlinear/pecos-models/"
mmap_model_path = f"/path/to/xlinear/mmap-models/"

print(f"Compiling mmap model from: {npz_model_path}, will save to : {mmap_model_path}...")
XLinearModel.compile_mmap_model(npz_model_path, mmap_model_path)
print("mmap model saved.")

Then user can load the memory-mapped model and do inference:

import sys
from pecos.xmc.xlinear.model import XLinearModel


mmap_model_path = f"/path/to/xlinear/mmap-models/"

# Load model
if sys.argv[2] == "--cmmap":
    print("Loading C/C++ mem map model...")
    xlm = XLinearModel.load(mmap_model_path, is_predict_only=True, is_mmap=True)
elif sys.argv[2] == "--cmmap-preload":
    print("Loading C/C++ mem map model pre-loaded...")
    xlm = XLinearModel.load(mmap_model_path, is_predict_only=True, is_mmap=True, pre_load=True)
else:
    print(f"Wrong option: {sys.argv[2]}")

# Load test data
Xt = XLinearModel.load_feature_matrix(f"/test/data/validation/X.npz")
Yt = XLinearModel.load_label_matrix(f"/test/data/validation/Y.npz")

# Predict
Yt_pred = xlm.predict(Xt)
Yt_pred = Yt_pred.tocsr()
metric = smat_util.Metrics.generate(Yt, Yt_pred, topk=10)
print(metric)

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

[rofuyu]

given that this variable is no longer a boolean, we might want to consider to remove the "b_" prefix.

[jiong-zhang]

Is compile_mmap_model implemented in HierarchicalMLModel? Also the namenpz_folder is kind of misleading.

[jiong-zhang]

Add docstrings for the new kwargs here and HierarchicalMLModel.load and XLinearModel.load

[jiong-zhang]

Also, can this be inferred from the saved configs rather than given by the user?

[jiong-zhang]

Should we add "compiled_format": "memory_map" in the params.json so that user need not to know the format before calling load?

[weiliw-amz]

Will break this PR into a series of PRs to merge:
#192
#189