multibrot_mojo_parallelize.mojo

from complex import ComplexSIMD
from math import iota
from algorithm import parallelize, vectorize
from tensor import Tensor
from utils.index import Index

alias float_type = DType.float64
alias simd_width = 2 * simdwidthof[float_type]()

alias width = 960
alias height = 960
alias MAX_ITERS = 200

alias min_x = -2.0
alias max_x = 0.6
alias min_y = -1.5
alias max_y = 1.5


fn mandelbrot_kernel_SIMD[
    simd_width: Int
](c: ComplexSIMD[float_type, simd_width]) -> SIMD[float_type, simd_width]:
    """A vectorized implementation of the inner mandelbrot computation."""
    let cx = c.re
    let cy = c.im
    var x = SIMD[float_type, simd_width](0)
    var y = SIMD[float_type, simd_width](0)
    var y2 = SIMD[float_type, simd_width](0)
    var iters = SIMD[float_type, simd_width](0)

    var t: SIMD[DType.bool, simd_width] = True
    for i in range(MAX_ITERS):
        if not t.reduce_or():
            break
        y2 = y * y
        y = x.fma(y + y, cy)
        t = x.fma(x, y2) <= 4
        x = x.fma(x, cx - y2)
        iters = t.select(iters + 1, iters)
    return iters


fn compute_multibrot_parallelized() -> Tensor[float_type]:
    let t = Tensor[float_type](height, width)

    @parameter
    fn worker(row: Int):
        let scale_x = (max_x - min_x) / width
        let scale_y = (max_y - min_y) / height

        @parameter
        fn compute_vector[simd_width: Int](col: Int):
            """Each time we operate on a `simd_width` vector of pixels."""
            let cx = min_x + (col + iota[float_type, simd_width]()) * scale_x
            let cy = min_y + row * scale_y
            let c = ComplexSIMD[float_type, simd_width](cx, cy)
            t.data().simd_store[simd_width](
                row * width + col, mandelbrot_kernel_SIMD[simd_width](c)
            )

        # Vectorize the call to compute_vector where call gets a chunk of pixels.
        vectorize[simd_width, compute_vector](width)

    # Parallelized
    parallelize[worker](height, height)
    return t


def main():
    _ = compute_multibrot_parallelized()

    # let multibrot = compute_multibrot_parallelized()
    # try:
    #     _ = show_plot(multibrot)
    # except e:
    #     print("failed to show plot:", e)


# def show_plot(tensor: Tensor[float_type]):
#     alias scale = 10
#     alias dpi = 64

#     from python import Python
#     np = Python.import_module("numpy")
#     plt = Python.import_module("matplotlib.pyplot")
#     colors = Python.import_module("matplotlib.colors")

#     numpy_array = np.zeros((height, width), np.float64)

#     for row in range(height):
#         for col in range(width):
#             numpy_array.itemset((col, row), tensor[col, row])

#     fig = plt.figure(1, [scale, scale * height // width], dpi)
#     ax = fig.add_axes((0.0, 0.0, 1.0, 1.0))
#     light = colors.LightSource(315, 10, 0, 1, 1, 0)

#     image = light.shade(numpy_array, plt.cm.hot, colors.PowerNorm(0.3), "hsv", 0, 0, 1.5)
#     plt.imshow(image)
#     plt.axis("off")
#     plt.savefig("multibrot_mojo_ parallelize.mojo.png")
#     plt.show()