High Performance Monte Carlo Simulation of Ising Model on TPU Clusters
Kun Yang Yi-Fan Chen Georgios Roumpos Chris Colby John Anderson





High Performance Monte Carlo Simulation of Ising Model on TPU Clusters
    Kun Yang, Yi-Fan Chen, George Roumpos, Chris Colby, John Anderson
    {kuny, yifanchen, roumposg, ccolby, janders}@google.com
======================================================================




Large-scale deep learning benefits from an emerging class of AI accelerators. Some of these accelerators' designs are general enough for compute-intensive applications beyond AI and Cloud TPU is one such example. In this paper, we demonstrate a novel approach using TensorFlow on Cloud TPU to simulate the two-dimensional Ising Model. TensorFlow and Cloud TPU framework enable the simple and readable code to express the complicated distributed algorithm without compromising the performance. Our code implementation fits into a small Jupyter Notebook and fully utilizes Cloud TPU's efficient matrix operation and dedicated high speed inter-chip connection. The performance is highly competitive: it outperforms the best published benchmarks to our knowledge by 60% in single-core and 250% in multi-core with good linear scaling. When compared to Tesla V100 GPU, the single-core performance maintains a ∼10% gain.  We also demonstrate that using low precision arithmetic—bfloat16—does not compromise the correctness of the simulation results.






§ INTRODUCTION

The Ising model <cit.>, which considers short-range interactions between spin variables on the sites of a d-dimensional lattice, plays an important role in statistical physics as a prototyping system to study the universal behavior of critical phenomena. Many significant breakthroughs in statistical physics are attributed to the study of the model from either its computational or its theoretical perspective. It is well known that the Ising model has no phase transition in one dimension; however, it undergoes a second-order phase transition between an ordered and a disordered phase in two dimensions or more <cit.>. The critical temperature T_c at which this phase transition occurs on a two-dimensional square lattice was analytically solved by Lars Onsager <cit.>, but it is still an open problem in three or more dimensions. Computer simulation offers a powerful alternative to study such systems and determine critical temperatures, thanks to the development of finite scaling theory <cit.> and availability of increasing computational power. This approach ushered in a plethora of interdisciplinary applications outside of physics, including bioinformatics <cit.>, economics <cit.> and operations research <cit.>.

Large-scale simulation of systems such as Ising model requires a large amount of high performance computing resources, which are usually available in multi-core computing architectures based on distributed shared memory, or distributed clusters (a.k.a data-centers) with homogeneous or heterogeneous nodes commonly seen in private or commercial clouds. Benefiting from the explosion of machine learning, especially deep learning, commercial clouds provide not only CPUs and GPUs, but also specialized chips such as FPGAs and other in-house processors. The Tensor Processing Unit (“Cloud TPU” or “TPU” for short)—an AI application-specific integrated circuit (ASIC) developed by Google for neural network machine learning—has received much attention in the machine learning community <cit.>. Its latest release, Cloud TPU v3,  offers 420× 10^12 floating-point operations per second (FLOPS) and 128GB of high bandwidth memory (HBM)[<cloud.google.com/tpu/>]. Multiple units are connected to form a “POD” (Cloud TPU v3 Pod) through a dedicated high speed 2-D toroidal mesh network, allowing up to 100+ peta-FLOPS and 32TB of HBM^<ref> to be accessed by the application with very low latency and in lockstep. TPU is programmable via software frontends such as TensorFlow <cit.> or PyTorch <cit.>, and can be deployed both for training huge deep neural networks and for performing low-latency online prediction. <cit.> reports impressive acceleration of training and online prediction.

With the tremendous amount of computation resources that TPU offers, it is compelling to also consider the opportunities TPU brings for applications beyond machine learning. The programming frontends that are used for TPU, such as TensorFlow, also offer a rich set of functionalities that are highly relevant for scientific computations. The TensorFlow TPU programming stack also provides the additional benefits of allowing distributed algorithms to be expressed with simple and easy-to-understand code without sacrificing performance. In addition, the ability to program conventional scientific simulations in TensorFlow framework makes it easier to explore the hybrid approaches employing both conventional scientific computation methods and modern machine learning techniques on the same framework.

Motivated by these observations, we developed a Single Instruction, Multiple Data (SIMD) distributed Markov Chain Monte Carlo (MCMC) simulation of the two-dimensional Ising model that is programmed in TensorFlow to run on TPU. We demonstrate that such an approach for scientific simulations is very promising. Our code implementation is easy to understand, with entire source code ∼600 lines, while also achieves competitive performance and scaling properties when compared with the state-of-the-art benchmarks in terms of both speed and scalability. Another interesting observation is that the lower precision arithmetic—bfloat16 (1 sign bit, 8 exponent bits and 7 mantissa bits [<en.wikipedia.org/wiki/Bfloat16_floating-point_format>]) instead of float32—does not compromise the accuracy of the result.

The implementation is open-sourced (URL available in AD Appendix) and can be run through Colaboratory[https://colab.research.google.com], a free cloud service based on Jupyter Notebooks for interactive data science. It is also worth pointing out that all results, including large-scale distributed computation, can be run through the Notebook frontend with minimal setup (mainly for allocating TPU backends and setting up the connection to the TPU backend).

In the following sections, we first review the high level architecture of TPU framework (Sec. 2), then discuss the adaption of a widely used algorithm for the Ising model simulation to better leverage TPU architecture (Sec. 3). We then follow up with benchmark results (Sec. 4) and in Section 5 highlight the implementation details and high-level performance analysis. Finally, we conclude the paper with Section 6, outlining our views on how the recent developments in both software and hardware for machine learning applications can impact the scientific computation applications.



§ TPU DEVICE ARCHITECTURE

TPU is a programmable linear algebra accelerator optimized for machine learning workloads. In the Cloud TPU v3 architecture, one TPU unit consists of four TPU chips on a board (we will use “TPU unit” or just “unit” to refer to such a board throughout this paper), and each TPU chip contains two TensorCores. Those TensorCores are treated as independent processors that communicate to each other through a dedicated high-bandwidth low-latency inter-chip network. In a larger system, more than 1000 TPU chips are packed on a two-dimensional toroidal mesh inter-chip interconnect network to form a TPU cluster known as a “TPU Pod” <cit.>.  In addition, each TPU unit is paired up with a TPU host server, with CPU, memory and disk resources attached. The communication links between TPU unit and its host server is through the regular data center connections. The host server can be used for data preparation, I/O task while TPU can be leveraged for compute-intensive tasks. The dedicated two-dimensional toroidal mesh allows all TensorCores in a “Pod” to work in lockstep efficiently, without going through the host servers. This large-scale high-bandwidth low-latency connectivity between TPU chips provides a significant advantage to achieve the strong linear-scaling performance reported in this paper.


The TensorCore, depicted in Figure <ref>, is optimized for dense linear algebra computations. It contains distinct classes of computing units, such as a scalar processor, a vector processor, accumulators and  matrix units <cit.>. All those processors are backed by its 16GB High-Bandwidth Memory (HBM). Vectorized operations are handled either in the vector processor directly or forwarded to corresponding extended vector units. Each extended vector unit takes the input operands, performs the corresponding operations, and returns the results back to the vector processor. One of these extended vector units is the matrix unit (MXU), which is capable of performing 128× 128 multiply-accumulate operations in each cycle <cit.>. The MXU is the main computing power of the TPU architecture, so it should be exploited as much as possible. While its inputs and outputs are 32-bit floating point values, the MXU rounds inputs down to bfloat16—a 16-bit floating point representation that provides better training and model accuracy than the IEEE half-precision representation—before multiplying .





Machine learning research shows that many machine learning models can tolerate lower precision arithmetic without degradation of converged accuracy. TPU supports storing values in bfloat16 format as a way to reduce the size of data and allow larger models to fit in memory <cit.>. For scientific computing, however, low precision is potentially dangerous because the increasing rounding error can introduce significant bias or divergent computation. In the case of the Ising model, while the binary spin values can be encoded in bfloat16 without loss of accuracy, the acceptance ratio and the random numbers used to determine acceptance in an MCMC simulation are more sensitive to reduced precision. However, our experiments show no noticeable differences in accuracy between bfloat16 and float32. By using bfloat16 instead of float32, we are able to simulate larger systems and leverage the MXU more effectively.

We program TPU through TensorFlow. The flow to run programs on TPU is roughly depicted in Figure <ref>. More details are available in TensorFlow's official XLA page[www.tensorFlow.org/xla]. In the first stage, TensorFlow constructs the computation graph and marks the graph for replication. Then the graph is rewritten to be TPU-compatible and compiled to a High Level Optimizer (HLO) program. Next, the Accelerated Linear Algebra (XLA) compiler takes over and converts HLO operations to Low Level Optimizer (LLO) code—effectively “TPU assembly code”, which can be readily executed on TPU. The graph construction and the compilation occurs on host server and incurs over-head cost. But once the compiled LLO code is deployed to TPU's the computation step can repeat as many times as required without the intervention from the host servers.   

Importantly XLA also provides communication primitives such as CollectivePermute[https://www.tensorflow.org/xla/operation_semantics#collectivepermute] and AllToAll[https://www.tensorflow.org/xla/operation_semantics#alltoall]. These primitives are implemented over the dedicated high-bandwidth low-latency inter-chip interconnect network so the communication between any TPU chips within a TPU Pod is extremely efficient. 

Another important detail that can have a critical impact on performance and memory usage  is the choice of the shape of the tensor variables used in the program (expressed as TensorFlow tensors). According to the performance guide <cit.>, unlike most other architectures, arrays in TPU are tiled in two dimensions. This entails to padding one dimension to a multiple of 8, and the other dimension to a multiple of 128. XLA performs data layout transformations and data are arranged in memory such that the hardware can efficiently process them. Programs that operate on array sizes undividable by 8 will have sub-optimal performance.



§ THE ISING MODEL

Mathematically, the Ising spin Hamiltonian is given by

    H(σ) = -J∑_⟨ ij⟩σ_iσ_j - μ∑_i=1^Nσ_i

where σ_i is a random variable assuming the values of ± 1 on sites i = 1, …, N of a d-dimensional hypercubic lattice, and ⟨ ij⟩ indicates that sites i and j are nearest neighbors. (Note that throughout this paper, we will use the bold σ to represent a tensor of spins, while the elements within it will be represented as regular σ). The first term, where the sum is over pairs of nearest-neighbor sites, represents the interaction energy that favors an ordered ferromagnetic state (if J > 0). The second term, involving the interaction between the applied field  and the spin system, is of a paramagnetic character. The configuration probability is given by the Boltzmann distribution with inverse temperature β = (k_BT)^-1:

    π(σ) = e^-β H(σ)/Z_β

where Z_β = ∑_σe^-β H(σ) is the partition function and k_B is the Boltzmann constant. For a function f of the spins (“observable”), we denote ⟨ f⟩=∑_σf(σ)π(σ) the expectation (mean value) of f. ⟨ f⟩ can often be difficult to evaluate numerically if there are many states in the system. The Markov Chain Monte Carlo is the most commonly used monte carlo algorithm to calculate statistics on the Ising model.

Without loss of generality, in what follows, we assume no external magnetic field, i.e., μ=0 and J=1, and the 2D lattice has circular boundary or in other words, is a torus. A given configuration of the lattice (spin values) is represented by matrix σ.



 §.§ Checkerboard Algorithm

The Metropolis-Hastings algorithm is the standard algorithm to simulate the Ising model. At each step, it proposes a candidate spin and flips the candidate based on the energy difference and acceptance probability. Closely related to this vanilla version that flips one spin at each step, there is another similar algorithm by flipping non-interacting spins in parallel, i.e., the efficient checkerboard algorithm <cit.>. 


Like the checkerboard above (Figure <ref>-(1)), spins in the lattice are colored black and white. The energy difference by flipping a spin of one color is completely described by its 4 neighbors of the opposite color. Thus, by fixing all spins of one color, the spins of the opposite color have no interactions with each other, and can be updated independently using Metropolis-Hastings. This observation leads to a highly parallel algorithm by alternating the 2 sub-routines below:



  * Fixing all black spins, flip each white spin based on Metropolis Hastings in parallel.

  * Fixing all white spins, flip each black spin based on Metropolis Hastings in parallel. 

Write π(σ) = π(σ_w, σ_b), where σ_w are values of all white spins and σ_b are values of all black ones. The transition kernel is,

    P{(σ_w, σ_b)→(σ_w^*, σ_b^*)}
       = P{(σ_w, σ_b)→(σ_w^*, σ_b)}P{(σ_w*, σ_b)→(σ_w^*, σ_b^*)}

and

    P{(σ_w, σ_b)→(σ_w^*, σ_b)}=∏_i∈wmin{1, e^β(σ_i^* - σ_i)·nn(i)}


    P{(σ_w^*, σ_b)→(σ_w^*, σ_b^*)}=∏_i∈bmin{1, e^β(σ_i^* - σ_i)·nn(i)}

where nn(i) is the sum of neighbor spins of i. By conditional probability decomposition, it is easy to show that

    π(σ_w^*, σ_b^*) = ∑_σ_w, σ_bP{(σ_w, σ_b)→(σ_w^*, σ_b^*)}π(σ_w, σ_b)

thus, the transition kernel satisfies the stationary distribution. The proof is based on Metropolis-Within-Gibbs Sampler <cit.> and is included in the supplemental materials. 



 §.§ Computation

The compute-intensive part of the checkerboard algorithm is the computation of the sum of neighbor values for each spin. To leverage the MXU, for a given lattice with size [128× m, 128× n] in a TPU core, we divide it into a [m, n] grid of 128× 128 sub-lattices (Figure <ref>-(1)). Define the kernel matrix K as,

    K = [ 0 1 0 … 0 0 0; 1 0 1 … 0 0 0; ⋮ ⋮ ⋮ ⋱ ⋮ ⋮ ⋮; 0 0 0 … 1 0 1; 0 0 0 … 0 1 0 ]_128× 128

then for a given sub-lattice σ_ij, matmul(σ_ij, K) + matmul(K, σ_ij) calculates the sum of nearest neighbors for all its internal sites. However, for the boundary sites, half of their nearest neighbors are missing from the sums. Those neighbors are on the boundaries of neighboring sub-lattices and ought to be sliced out and added into the sums. The whole lattice is updated one color at a time. To fix the spins of one color on the checkerboard, we multiply flip probabilities with a binary mask M, where

    M = [ 1 0 1 … 0 1 0; 0 1 0 … 1 0 1; ⋮ ⋮ ⋮ ⋱ ⋮ ⋮ ⋮; 1 0 1 … 0 1 0; 0 1 0 … 1 0 1 ]_128× 128

In summary, the algorithm to update the lattice given the color c and configuration σ is given in Algorithm <ref>. By alternating colors with different c, the Ising model is properly simulated. However, in Algorithm <ref>, there are several redundant calculations that are subject to further optimization: 1) Line 1 generates the probability for all spins, while only the spins colored by c are eligible for flipping. 2) Lines 2-6 calculate the nearest neighbor sums of all spins, however, only spins colored by c are updated. 3) Lines 8-9 generate flips but multiplying mask to fix the opposite color is expensive. To eliminate all the redundancies above, we reorganize the lattice and represent it in a compact way, i.e., the lattice is instead split into a [m', n'] grid of [256, 256] sub-lattices, and define 4 “compact” [m', n', 128, 128] sub-lattices (Figure <ref>-(2)),

    σ̂_00 = σ[:, :, 0::2, 0::2], σ̂_01 = σ[:, :, 0::2, 1::2]


    σ̂_10 = σ[:, :, 1::2, 0::2], σ̂_11 = σ[:, :, 1::2, 1::2]

and kernel,

    K̂ = [ 1 1 0 … 0 0 0; 0 1 1 … 0 0 0; ⋮ ⋮ ⋮ ⋱ ⋮ ⋮ ⋮; 0 0 0 … 0 1 1; 0 0 0 … 0 0 1 ]_128× 128

thus, σ̂_00 and σ̂_11 are all “black” spins, and σ̂_01 and σ̂_10 are all “white” spins. Now, it is trivial to show that the nearest neighbor sums of internal spins of those 4 “compact” sub-lattices are

    nn(σ̂_00)=matmul(σ̂_01, K̂) + matmul(K̂^T, σ̂_10) 
       nn(σ̂_11)=matmul(K̂, σ̂_01) + matmul(σ̂_10, K̂^T) 
       nn(σ̂_01)=matmul(σ̂_00, K̂^T) + matmul(K̂^T, σ̂_11) 
       nn(σ̂_10)=matmul(K̂, σ̂_00) + matmul(σ̂_11, K̂)

and their boundary spins are corrected using a similar approach as in Algorithm <ref>. The complete algorithm is presented in Algorithm <ref>. According to our experiments, it is about 3x faster than Algorithm <ref> and has less memory footprint as it uses less temporary HBM.









§ SIMULATION RESULTS



 §.§ Correctness

The average magnetization per spin for a given β is defined as,

    m(T) = m(β) =⟨σ⟩=1/N∑_iσ_i

where N = n^2 for a square lattice with size n. It can be shown that below the critical temperature, i.e.,
(k_Bβ)^-1 = T < T_c = (k_Bβ_c)^-1=2/k_Bln(1 + √(2)), there is spontaneous magnetization—the interaction among spins is sufficiently large to cause neighboring spins to spontaneously align. On the other hand, thermal fluctuations completely eliminate any alignment above the critical temperature. Moreover, at the critical temperature, there is a discontinuity in the first derivative of ⟨ m(T)⟩ with respect to T_c. This discontinuity generates a downward drop in the average magnetization. The sudden loss of spontaneous magnetization above the critical temperature is a signature of phase transition. Besides average magnetization, a more sensitive test of correctness is the Binder parameter <cit.>, which is given by

    U_4(T) = 1 - ⟨ m(T)^4⟩/3⟨ m(T)^2⟩^2

namely, the kurtosis of m(T). It is frequently used to accurately determine phase transition points in numerical simulations of various models.

We verify the correctness of our algorithm and implementation by computing both quantities at various temperatures on different sizes of square lattices. As shown in Figure <ref>, the curves of average magnetization, with subtle differences, overlap with each other, and those of Binder parameters cross the critical line almost perfectly. Additionally, we investigate the implications of lower precision, i.e., bfloat16, for the estimation accuracy. In MCMC simulation of Ising model, lower precision has impact on the calculation of acceptance ratio and the random number generation, the bias introduced might accumulate and decrease the overall accuracy. However, in our experiments, all curves generated in bfloat16 and float32, especially those of Binder parameters, are almost the same and sharp turns around critical line are clearly observable. Based on this evidence, we argue that using bfloat16 has negligible impact on Ising model simulation, and in turn it offers two benefits: 1). We are able to simulate a larger lattice on a single TPU core because of smaller memory footprint, and 2). bfloat16 matrix multiplication with 32-bit accumulation is very efficient in MXU, while float32 matrix multiplication is more expensive as several bfloat16 passes are required.





 §.§ Benchmarks

Preis et al. <cit.> showed impressive acceleration in their single GPU implementation of checkerboard algorithm using the compute unified device architecture (CUDA). By exploiting GPU's  large pool of threads, its single-instruction multiple thread (SIMT) unit and memory hierarchy, their algorithm achieved 60x speedup on single GPU compared to its CPU counterpart. In their follow-up paper <cit.>, their original algorithm was modified to overcome the memory limitations of single GPU. The improved algorithm was able to simulate significantly larger systems and reached a performance of ∼7.9774 flips/ns throughput in its best performing variant. By combining CUDA with MPI on the CPU level, their distributed algorithm achieved 206 flips/ns on a 800,000^2 lattice. Besides the work on GPU, another encouraging line of research is the use of field-programmable gate array (FPGA). A recent implementation achieves ∼ 614.4 flips/ns throughput, see <cit.> and references therein.

To quantify the performance of our implementation, we run our algorithm on TPU v3 using single core and multiple cores on TPU v3 clusters. As in <cit.>, we measure the time spent on one sweep update, i.e., one update on all “black” spins plus one update on all “white” ones, and compute the average number of flips per nanosecond by dividing with n^2.


  §.§.§ Single TPU Core

First, we simulate the system on a single TPU v3 core (half TPU v3 chip). Our choice of lattice size is compatible with MXU's registers and achieves 100% memory capacity utilization according to our profilings. We can simulate lattice  with size up to (656× 128)^2 = 83,968^2, which consumes 96% of the memory. The benchmarks in various sizes from (20× 128)^2 to (640× 128)^2 are summarized in Table. <ref>, the lattice size and flips in nanoseconds increase in tandem, as more computation is spent on matrix multiplication. For comparison, we also implemented the algorithms based on <cit.> under , and using its  package for random number generation and its  package for reductions. To avoid the excessive temporary memory allocation on GPU, we wrote a custom memory allocator to reuse temporary memory. The benchmark we obtained on NVIDIA's Tesla V100, which powered by NVIDIA's latest Volta architecture, is 11.3704 flips/ns. Other older benchmarks published in <cit.> are also listed for reference. 



It is also interesting to estimate the energy consumption in the computation. Assuming the average power consumption during the operation to be 𝒫W and the throughput achieved to be ℱflips/ns, the corresponding energy used is (𝒫/ℱ)nJ/flip. The actual average power consumption depends on many factors and usually sophisticated modeling and measurements are needed.
However, for a rough estimate of the upper bound, we use 250 W for GPU (based on NVIDIA's Tesla V100 Spec for PCIe version max power consumption[<images.nvidia.com/content/technologies/volta/pdf/437317-Volta-V100-DS-NV-US-WEB.pdf>]). 
While Google doesn't publish the number for TPU v3, it has been estimated to be 200W[<www.nextplatform.com/2018/05/10/tearing-apart-googles-tpu-3-0-ai-coprocessor/>] for a TPU v3 chip, or equivalently 100W for a TPU v3 core.



  §.§.§ Linear Scaling on Multiple TPU cores

The idea to simulate Ising model on TPU v3 clusters is to split the whole lattice into sub-lattices, and exchange their boundary values to calculate the nearest-neighbor sums and update the sub-lattices in each core independently. The challenges are to handle the synchronization among the cores: 1). Block the update of sub-lattices until all the nearest-neighbor sums are calculated. 2). Block next iteration before the update of all sub-lattices is completed. Fortunately, such synchronization is already implemented in TensorFlow op , which is used to exchange data by specifying source and target cores. 

In its current release, TPUs in a TPU Pod are organized into a grid and each TPU core has an associated coordinate. All cores communicate through a specialized high-speed interconnect. In our experiments, we use smaller sections of a pod called slices <cit.>, and show that the TPU Pod interconnect makes the overhead of exchanging boundary values between cores negligible. As a result, we observe strict linear scaling of flips/ns to the number of TPU cores, as shown in Table. <ref>.

In <cit.>, the communications between GPUs are handled by MPI through the hosts, which is potentially a bottleneck in the simulations. Another interesting comparison is to benchmark their multi-GPU algorithm using , which enables the interconnect of 8 GPUs, or  that support the interconnect of 16 GPUs <cit.>. However, as the code for multi-GPU simulation is not available and the limitations on the number of interconnect GPUs, we leave it for our future work.





§ IMPLEMENTATION HIGHLIGHTS AND PERFORMANCE ANALYSIS

In the previous section, we report the competitive performance numbers achieved in our experiments. It is worthwhile going into some depth to highlight the implementation design choices and analyze the performance, in order to provide more insights on this new approach towards scientific simulations. The entire source code is also available through AD appendix.    


 §.§ Implementation Highlights: TensorFlow, SIMD, Highspeed Mesh Network

While most commonly seen applications of TensorFlow are in the context of machine learning, the functionalities available are highly relevant to various scientific applications. At a high level, TensorFlow provides various transformation operations that can be applied on tensors. These tensors can be either stateful variables or temporary values. 

It is easy to see that the two-dimensional spin sites can be easily represented as a stateful tensor variable. In this work, we have chosen to use arrays of rank-4 tensors to represent the supergrid structure used in checkerboard algorithm  <ref>. More specifically, the partial spin lattice on each core is represented as a four-dimensional array: super_grids[N_x, N_y, 2, 2] where each of the array elements is a rank-4 tensor variable with shape of [m, n, 128 × i, 128 × j]. 

The tensor shape is chosen so the last two dimensions are always integer multiples of 128, to better match TPU HBM tiling and the MXU structure<cit.>. We also choose to use super_grids[:, :, 0, 0], super_grids[:, :, 1, 1] to represent the black compact sub-lattices, and use super_grids[:, :, 0, 1] and super_grids[:, :, 1, 0] to represent the white compact sub-lattices, as depicted in Figure <ref>-(2). In the results reported in Table <ref> we use (N_x, N_y, m, n, i, j) = (2, 2, 224, 112, 1, 1), which gives us the per-core lattice size of 128×[896, 448].

To determine the acceptance for the flipping of each spin site, a  random tensor generation operation available in TensorFlow is used: tf.random_uniform[<www.tensorflow.org/api_docs/python/tf/random/uniform>].
It generates random tensor for a given shape with uniform probability between [0, 1] for each element. While this process is not the most compute-intensive, it does take up about ∼10% of the step time (more discussions on this in next sub-section).

The most compute-intensive part of the simulation is the computation of acceptance ratio, which involves summing on the nearest neighbor spin values. As pointed out previously, we leverage MXU's matrix multiplication to achieve this. Since each TensorCore provides a raw computational power of ∼50 TFLOPS, it greatly helps the efficiency of our simulation. This part takes up about ∼60% of the step time.

In TensorFlow framework, the expression of the computation is converted into a graph, and for TPU, it is further compiled through XLA into the executable programs (LLO) and deployed to TPU during run time. This Just-In-Time (JIT) compilation can incur overhead but usually it is fairly small (under a few secs) for smaller problems, and while it can sometimes take longer time (up to minutes) for larger problems, usually it is well-amortized for these larger problems as typically millions of steps are executed.

For the distributed multi-core case, TensorFlow also provides the primitive for SIMD programming on TPU. Using tf.tpu.replicate[<www.tensorflow.org/api_docs/python/tf/tpu/replicate>], one can easily replicate the computation across multiple TPU TensorCores with simple syntax (typically a few lines of code) to map and assign core ids using an object returned during the TPU system initialization call, which encapsulates the mesh topology information.  



Another critical component for the distributed case is that in the acceptance ratio computation, each core needs to exchange its border spin values with the neighboring cores. Again, TensorFlow provides a simple primitive that is syntactically simple and leverages TPU Pod's high speed dedicated inter-chip mesh network, tpu_ops.collective_permute[<www.github.com/tensorflow/tensorflow/blob/master/tensorflow/python/tpu/ops/tpu_ops.py>] which allows a tensor's value to be permuted through different cores according to the source-destination pair specification. Note that in this case, each core has the same instruction and in the call to tpu_ops.collective_permute, the source-destination mapping contains globally identical specifications. Each core that is involved in the operation, when executing this part of program, will block until it sends and receives the corresponding values according to the source-destination specification. 

Figure <ref> shows an example of a collection of 3 cores that are exchanging the 'boundaries' with periodic boundary condition. Each core has tensors R and L representing the right and left boundaries of the internal grid. Figure <ref> shows how each core can acquire the extended boundaries from its neighboring cores. The highly efficient inter-core communications in TPU Pod allows all TensorCores to work in lockstep with minimal latency even when a large number of cores participate in the communications. In fact, the time spent on this step in our experiments is well below 0.15% in all cases and this explains the linear-scaling performance we observed. Note that in all our experiments, no attempt is made to match the logical layout of the lattice with the physical TPU cores: two logically neighboring sub-lattices could be distributed to two cores physically far apart, requiring multiple hops for the communication between them and, yet this doesn't cause any noticeable performance loss.      



 §.§ Performance Analysis


We look into several key components of the computations and their respective performance in terms of the time spent on each. We use the tool devloped for TPU profiling (available also in Google Cloud[<cloud.google.com/tpu/docs/cloud-tpu-tools>]. The profiling tool is able to provide fine-grained analysis of the utilization of the hardware, the efficiency of the operations at program level and more. Figure <ref> gives an example output from the tool's trace viewer. 


Using the tool, we took measurements of the breakdown of the key operations at the HLO level: the time spent on the computation of the nearest neighbor sum (mostly using MXU), the time spent on generating the random uniform tensors (mostly using VPU), the time spent of data formatting and reshaping for the computation, and the time spent on data exchange between cores. Detailed breakdown is in Table <ref>. Note that from Table <ref> we know that for all cases, the time step is essentially identical (∼580ms), so percentage numbers in this table can be used directly for comparison across different cases. 


The first observation is that the time breakdown is fairly stable across different scales. We also see that the most expensive part is the computation of nearest neighbor sum, leveraging the MXU, which accounts for ∼60% of the time. The time spent on generating random tensors is also significant: it is  ∼12% of the time. The data formatting takes up ∼28% of the time. It is also worth noting that this part can be significantly worse if the shape of the tensor variables do not conform to the tiling in TPU HBM. The most interesting part is the time spent on inter-core communication, and it can be seen that it is taking very insignificant amount of time, even for the very large case when 512 cores are involved. It is a key property of TPU Pod that allows the linear scaling. Note that since the cores participating in the “collective_permute” operations will perform both sending and receiving and the time measured includes both the synchronization overhead between the cores as well as the time that the data travels between cores. However, since the data amount (the edges of the sub-lattice) is small, the time is not dominated by the data propagation and not bandwidth bound (the largest edge is only 229,376 bytes, and would take just ∼0.023msec to transport even on a relatively moderate 10GB network). Table  <ref> shows the step time and the “collective_permute” time with various per-core lattice size at different number of cores. The “collective_permute” time in all cases are insignificant when compared with the step time. We also note that this time is more directly affected by the number of cores than the size of the sub-lattice, indicating that the communication is not in the bandwidth limited regime. 

Another interesting aspect of the data in Table <ref> is that, for a constant full lattice size (i.e., the entries along the diagonal direction in Table <ref>), while the step time consistently decreases as the number of cores increases, we can see that there are two different regimes of the rate of step time decrease. When the per-core lattice size decreases from [896×128, 448×128] to [448×128, 224×128], a 4× decrease, the step time decreases from ∼575 ms to ∼255ms, or ∼44%, instead of 25%. But when the per-core lattice size decrease from [896×128, 448×128] to [224×128, 112×128], the step time changes from ∼255 ms to ∼65ms, or ∼25.5%. This is due to higher MXU utilization when the per-core lattice size is [896×128, 448×128], and the utilization pattern changes (lower MXU utilization) when the per-core lattice size changes to [448×128, 224×128], and stays about the same when the per-core lattice size is [224×128, 112×128]. 



Finally, we use the profiling tool to measure the FLOPS performance of our program. The results are summarized in Table <ref>. In all cases, the achieved FLOPS is about 76% of the roofline (memory bound), or roughly 5.89 TFLOPS. A rough estimate can be done using the number operations in matrix multiplication per core used for nearest neighbor sum: there are total 4×896×448 of matrix multiplications of size 128  128 (2× from the inner grid nearest neighbor computations, and 2× from the boundary nearest neighbor). Total number of operations is 896×448×128^3. Using step time ∼580ms, we got the estimate of 5.8 TFLOPS, vey close to the measured program FLOPS. It is also interesting to note that, from the slope of the roofline plot in Table <ref>, we can estimate the HBM bandwidth to be at least ∼300GB/sec.

We believe that the efficiency of the computation can be improved by further optimizing the data formatting operations: identifying bottlenecks and rearranging the layout of the tensors. It is also worth noting that the matrix multiplication involves sparse diagonal band kernel with shape of 128×128 and we could potentially explore smaller size of kernel to improve efficiency. A possible direction is trying to utilize the convolution operation in TensorFlow to improve the efficiency further. We also expect that as XLA being actively developed and improved over time, it will deliver higher and higher performance.  



§ CONCLUSION

We demonstrate a novel approach to simulate the two-dimensional ferromagnetic Ising model using TensorFlow on Cloud TPU. We adapted the standard checkerboard algorithm to exploit the TPU architecture, in particularly to leverage its efficient matrix operation and the dedicated high-bandwidth low-latency inter-chip interconnect network of the TPU Pod. We calculate the average magnetization and Binder parameter at various temperatures with different lattice sizes. Our numeric estimates on those size-independent quantities match the theoretical results using both float32 and bfloat16 precision. Our benchmarks also demonstrate competitive and linear-scaling performance. The algorithm used in this work can be generalized for three-dimensional Ising model. An interesting direction to follow up would be applying the approach on some of the recent works that push the three-dimensional Ising model simulations to limits, for example<cit.>.


This work demonstrates how the new Cloud TPU computation resources could be efficiently employed for conventional scientific simulation problems. However, even more significantly, by implementing the entire simulation using TensorFlow framework, we point a direction where the direct integration of machine learning approaches with conventional simulations is possible and could be done easily. For example, the automatic differentiation in TensorFlow[<www.tensorflow.org/tutorials/eager/automatic_differentiation>] is readily applicable for optimization of parameters in design problems that use simulations. In the context of this current work, an interesting followup would be finding the optimal J_i,j given material properties for the case where J is not uniform across all spin sites. In our view, the research in this direction will bring many interesting advancements and will continue to shape the future of computing.  



§ APPENDICES



 §.§ Proof of Stationarity


    ∑_σ_w, σ_bP{(σ_w, σ_b)→(σ_w^*, σ_b^*)}π(σ_w, σ_b) 
    
    =   ∑_σ_w, σ_bP{(σ_w, σ_b)→(σ_w^*, σ_b)}P{(σ_w^*, σ_b)→(σ_w^*, σ_b^*)}π(σ_w, σ_b) 
    
    =   ∑_σ_w, σ_b∏_i∈ wP(σ_i→σ_i^*|σ_b)∏_i∈ bP(σ_i→σ_i^*|σ_w^*)∏_i∈ wπ(σ_i|σ_b)·π(σ_b) 
    
    =   ∑_σ_w, σ_b∏_i∈ wP(σ_i→σ_i^*|σ_b)π(σ_i|σ_b)∏_i∈ bP(σ_i→σ_i^*|σ_w^*)·π(σ_b) 
       By detailed balance of single spin flip:
        P(σ_i→σ_i^*|σ_b)π(σ_i|σ_b) = P(σ_i^*→σ_i|σ_b)π(σ_i^*|σ_b) 
      
    =   ∑_σ_w, σ_b∏_i∈ wP(σ_i^*→σ_i|σ_b)π(σ_i^*|σ_b)∏_i∈ bP(σ_i→σ_i^*|σ_w^*)·π(σ_b) 
    
    =   ∑_σ_w, σ_b∏_i∈ wP(σ_i^*→σ_i|σ_b)∏_i∈ bP(σ_i→σ_i^*|σ_w^*)∏_i∈ wπ(σ_i^*|σ_b)·π(σ_b) 
    
    =   ∑_σ_w, σ_b∏_i∈ wP(σ_i^*→σ_i|σ_b)∏_i∈ bP(σ_i→σ_i^*|σ_w^*)π(σ_i|σ_w^*)·π(σ_w^*) 
       By detailed balance of single spin flip:
        P(σ_i→σ_i^*|σ_w^*)π(σ_i|σ_w^*)=P(σ_i^*→σ_i|σ_w^*)π(σ_i^*|σ_w^*) 
    
    =   ∑_σ_w, σ_b∏_i∈ wP(σ_i^*→σ_i|σ_b)∏_i∈ bP(σ_i^*→σ_i|σ_w^*)π(σ_i^*|σ_w^*)·π(σ_w^*) 
    
    =   π(σ_w^*, σ_b^*)∑_σ_w, σ_b∏_i∈ wP(σ_i^*→σ_i|σ_b)∏_i∈ bP(σ_i^*→σ_i|σ_w^*) 
    
    =   π(σ_w^*, σ_b^*)















 §.§ Further Optimization and Scaling

In this section, we present additional optimization and further up-scaling of the problem. The additional optimization involves a detailed implementation of the computation of the nearest neighbor energy contribution:  is used, instead of batch multiplication. Detailed implementation is also open-sourced and can be found in [ ]. This approach more efficiently leverages MXU's computation power by packing more operations together for each memory load operation. Together with the improvements from the new version of TensorFlow (r1.15), we achieve a ∼80% performance improvement. In addition, we also utilize all available 2048 TPU cores in a TPU v3 POD (previously, we only utilized one quarter of a full pod). 

Figure <ref> shows the simulated magnetization and Binder parameters using the new implementation and it confirms the new algorithm continues to produce the correct results. 

In our weak-scaling performance tests, we explored different density of work-load: loose-packed ([224, 224] × 128 per core), dense-packed ([448, 448] × 128 per core) and superdense-packed ([896, 448] × 128 per core). We demonstrated that in all cases, they all scale linearly and the largest problem we can handle is 4
× of the largest problem we previously reported. Table <ref> provides detailed results. We also provide a plot of all available reported performance numbers  in Figure <ref>.

We also perform strong-scaling performance tests. The size of the problem we choose is the (128 × 1792)^2. Table <ref> shows the results. The scaling stays relatively linear for smaller number of cores, but when more than 1000 cores are involved, the overhead of communication starts to be a significant part of the run time. From Figure <ref> this can also be observed clearly







§ ACKNOWLEDGEMENTS

We would like to thank Blake Hechtman, Brian Patton, Cliff Young, David Patterson, Naveen Kumar, Norm Jouppi, Rif A. Saurous, Yunxing Dai, Zak Stone, and many more for valuable discussions and helpful comments, which have greatly improved the paper.



SUPPLEMENTAL MATERIALS


Colaboratory is hosted on  in the directory of .

plain
10

[Abadi(2016)]Abadi2016
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey
  Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, and Michael Isard.
Tensorflow: a system for large-scale machine learning.
In OSDI, volume 16, pages 265–283, 2016.

[Pushing-To-Limit-3D(2018)]Pushing-To-Limit-3D-2018
Jiahao Xu Alan M. Ferrenberg and David P. Landau.
Pushing the limits of monte carlo simulations for the 3d ising model.
arXiv preprint arXiv:1806.03558, 2018.

[Block(2010)]Block2010
Peter Virnau Benjamin Block and Tobias Preis.
Multi-gpu accelerated multi-spin monte carlo simulations of the 2d
  ising model.
Computer Physics Communications, 181(9):1549–1556, 2010.

[Binder(1981)]Binder1981
Kurt Binder.
Finite size scaling analysis of ising model block distribution
  functions.
Zeitschrift für Physik B Condensed Matter, 43(2):119–140,
  1981.

[Binder(2001)]Binder2001
Kurt Binder and Erik Luijten.
Monte carlo tests of renormalization-group predictions for critical
  phenomena in ising models.
Physics Reports, 344(4-6):179–253, 2001.

[Cook(2018)]Cook2018
Takashi Sato Masayuki Hiromoto Chase Cook, Hengyang Zhao and Sheldon X.-D. Tan.
Gpu based parallel ising computing for combinatorial optimization
  problems in vlsi physical design.
arXiv preprint arXiv:1807.10750, 2018.

[tpupod(2019)]tpupod2019
Google Cloud.
Choosing between a single cloud tpu device and a cloud tpu pod
  (alpha).
https://cloud.google.com/tpu/docs/deciding-pod-versus-tpu,
  2019.

[tpuref(2019)]tpuperf2019
Google Cloud.
Performance guide.
https://cloud.google.com/tpu/docs/performance-guide, 2019.

[tpusys(2019)]tpusys2019
Google Cloud.
System architecture.
https://cloud.google.com/tpu/docs/system-architecture, 2019.

[tpubfloat16(2019)]tpubfloat16
Google Cloud.
Using bfloat16 with tensorflow models.
https://cloud.google.com/tpu/docs/bfloat16, 2019.

[Barradas-Bautista(2018)]Barradas-Bautista2018
Mark Agostino Didier Barradas-Bautista, Matias Alvarado-Mentado and Germinal
  Cocho.
Cancer growth and metastasis as a metaphor of go gaming: An ising
  model approach.
PloS one, 13(5):e0195654, 2018.

[nvlink(2019)]nvlink2019
Nvidia NVLink Fabric.
Nvlink fabric: Advancing multi-gpu processing.
https://www.nvidia.com/en-us/data-center/nvlink/, 2019.

[Ising(1925)]Ising1925
Ernst Ising.
Beitrag zur theorie des ferromagnetismus.
Zeitschrift für Physik, 31(1):253–258, 1925.

[Jouppi(2017a)]Jouppi2017a
Norm Jouppi.
Quantifying the performance of the tpu, our first machine learning
  chip.
Google Cloud:
  https://cloud.google.com/blog/products/gcp/quantifying-the-performance-of-the-tpu-our-first-machine-learning-chip,
  2017.

[Binder(1993)]Binder1993
Lyle Roelofs A. John Mallinckrodt Kurt Binder, Dieter Heermann and Susan McKay.
Monte carlo simulation in statistical physics.
Computers in Physics, 7(2):156–157, 1993.

[Mueller(1991)]Mueller1991
Peter Müller.
A generic approach to posterior integration and Gibbs sampling.
Purdue University, Department of Statistics, 1991.

[Jouppi(2019)]Jouppi2019
George Kurian Sheng Li Nishant Patil James Landon Cliff Young Norman P. Jouppi,
  Doe Hyun Yoon and David Patterson.
A domain-specific supercomputer for training deep neural networks.
Submitted to Communications of the ACM, 2019, 2019.

[Jouppi(2017)]Jouppi2017
Nishant Patil David Patterson Gauray Agrawal Raminder Bajwa Sarah Bates Suresh
  Bhatia Nan Boden Norman P. Jouppi, Cliff Young and Al Borchers.
In-datacenter performance analysis of a tensor processing unit.
In Computer Architecture (ISCA), 2017 ACM/IEEE 44th Annual
  International Symposium on, pages 1–12. IEEE, 2017.

[Onsager(1994)]Onsager1944
Lars Onsager.
Crystal statistics. i. a two-dimensional model with an order-disorder
  transition.
Physical Review, 65(3-4):117, 1944.

[Ortega-Zamorano(2016)]Ortega-Zamorano2016
Francisco Ortega-Zamorano, Marcelo A. Montemurro, Sergio Alejandro Cannas,
  José M. Jerez, and Leonardo Franco.
Fpga hardware acceleration of monte carlo simulations for the ising
  model.
IEEE Trans. Parallel Distrib. Syst., 27(9):2618–2627, 2016.

[Paszke(2017)]Paszke2017
Adam Paszke, Sam Gross, Soumith Chintala, and Gregory Chanan.
Pytorch: Tensors and dynamic neural networks in python with strong
  gpu acceleration.
PyTorch: Tensors and dynamic neural networks in Python with
  strong GPU acceleration, 2017.

[Preis(2008)]Preis2008
Tobias Preis, Wolfgang Paul, and Johannes J. Schneider.
Fluctuation patterns in high-frequency financial asset returns.
EPL (Europhysics Letters), 82(6):68005, 2008.

[Preis(2009)]Preis2009
Tobias Preis, Peter Virnau, Wolfgang Paul, and Johannes J. Schneider.
Gpu accelerated monte carlo simulation of the 2d and 3d ising model.
Journal of Computational Physics, 228(12):4468–4477, 2009.

[Prieto-Castrillo(2018)]Prieto-Castrillo2018
Francisco Prieto-Castrillo, Amin Shokri Gazafroudi, Javier Prieto, and
  Juan Manuel Corchado.
An ising spin-based model to explore efficient flexibility in
  distributed power systems.
Complexity, 2018, 2018.

[Romero(2019)]Romero2019
Joshua Romero, Mauro Bisson, Massimiliano Fatica, and Massimo Bernaschi.
A performance study of the 2D ising model on GPUs.
arXiv preprint arXiv:1906.06297, 2019, 2019.



