A primary concern in creating compartments (i.e., protection domains) for bare-metal systems is to adopt the applicable compartmentalization policy. Existing studies have proposed several typical policies in literature. However, neither of the policies consider the influence of unsafe functions on the compartment security that a vulnerable function would expose unpredictable attack surfaces, which could be exploited to manipulate any contents that are stored in the same compartment. In this paper, we design a machine learning-assisted compartmentalization scheme, which adopts a new policy that takes every function’s security into full account, to create compartments for bare-metal systems. First, the scheme takes advantage of the machine learning method to predict how likely a function holds an exploitable security bug. Second, the prediction results are used to create a new instrumented firmware that isolates vulnerable and normal functions into different compartments. Further, the scheme provides some optional optimization plans to the developer to improve the performance. The PoC of the scheme is incorporated into an LLVM-based compiler and evaluated on a Cortex-M based IoT device. Compared with the firmware adopting other typical policies, the firmware with the new policy not only shows better security but also assures the overhead basically unchanged.