Jie Wang

Jian Weng, Animesh Jain, Jie Wang et al.

Because of the increasing demand for computation in DNN, researchers develope both hardware and software mechanisms to reduce the compute and memory burden. A widely adopted approach is to use mixed precision data types. However, it is hard to leverage mixed precision without hardware support because of the overhead of data casting. Hardware vendors offer tensorized instructions for mixed-precision tensor operations, like Intel VNNI, Tensor Core, and ARM-DOT. These instructions involve a computing idiom that reduces multiple low precision elements into one high precision element. The lack of compilation techniques for this makes it hard to utilize these instructions: Using vendor-provided libraries for computationally-intensive kernels is inflexible and prevents further optimizations, and manually writing hardware intrinsics is error-prone and difficult for programmers. Some prior works address this problem by creating compilers for each instruction. This requires excessive effort when it comes to many tensorized instructions. In this work, we develop a compiler framework to unify the compilation for these instructions -- a unified semantics abstraction eases the integration of new instructions, and reuses the analysis and transformations. Tensorized instructions from different platforms can be compiled via UNIT with moderate effort for favorable performance. Given a tensorized instruction and a tensor operation, UNIT automatically detects the applicability, transforms the loop organization of the operation,and rewrites the loop body to leverage the tensorized instruction. According to our evaluation, UNIT can target various mainstream hardware platforms. The generated end-to-end inference model achieves 1.3x speedup over Intel oneDNN on an x86 CPU, 1.75x speedup over Nvidia cuDNN on an NvidiaGPU, and 1.13x speedup over a carefully tuned TVM solution for ARM DOT on an ARM CPU.

Tian-Xiang Mao, Jie Wang, Baojiu Li et al.

We propose a new scheme to reconstruct the baryon acoustic oscillations (BAO) signal, which contains key cosmological information, based on deep convolutional neural networks (CNN). Trained with almost no fine-tuning, the network can recover large-scale modes accurately in the test set: the correlation coefficient between the true and reconstructed initial conditions reaches $90\%$ at $k\leq 0.2 h\mathrm{Mpc}^{-1}$, which can lead to significant improvements of the BAO signal-to-noise ratio down to $k\simeq0.4h\mathrm{Mpc}^{-1}$. Since this new scheme is based on the configuration-space density field in sub-boxes, it is local and less affected by survey boundaries than the standard reconstruction method, as our tests confirm. We find that the network trained in one cosmology is able to reconstruct BAO peaks in the others, i.e. recovering information lost to non-linearity independent of cosmology. The accuracy of recovered BAO peak positions is far less than that caused by the difference in the cosmology models for training and testing, suggesting that different models can be distinguished efficiently in our scheme. It is very promising that Our scheme provides a different new way to extract the cosmological information from the ongoing and future large galaxy surveys.