starpu-max/julia/mult/mult.jl

import Libdl
using StarPU
using LinearAlgebra

#shoud be the same as in the makefile
const STRIDE = 72

@target STARPU_CPU+STARPU_CUDA
@codelet function matrix_mult(m1 :: Matrix{Float32}, m2 :: Matrix{Float32}, m3 :: Matrix{Float32}) :: Float32

    width_m2 :: Int32 = width(m2)
    height_m1 :: Int32 = height(m1)
    width_m1 :: Int32 = width(m1)
    # Naive version
    #@parallel for j in (1 : width_m2)
    #    @parallel for i in (1 : height_m1)
    #
    #          sum :: Float32 = 0.

    #          for k in (1 : width_m1)
    #              sum = sum + m1[i, k] * m2[k, j]
    #          end
    
    #          m3[i, j] = sum
    #      end
    #  end
    ##### Tiled and unrolled version 
    for l in (1 : width_m2)
        for m in (1 : height_m1)
            m3[m,l] = 0
        end
    end
    @parallel for i in (1 : STRIDE : height_m1)
        for k in (1 : STRIDE : width_m1 )
            for j in (1 : STRIDE : width_m2  )
                for kk in (k : 4 : k+STRIDE-1)
                    for jj in (j : 2 : j+STRIDE-1)
                        alpha00 :: Float32 =m2[kk,jj]
                        alpha01 :: Float32 =m2[kk,jj+1]
                        alpha10 :: Float32 =m2[kk+1,jj]
                        alpha11 :: Float32 =m2[kk+1,jj+1]
                        alpha20 :: Float32 =m2[kk+2,jj]
                        alpha21 :: Float32 =m2[kk+2,jj+1]
                        alpha30 :: Float32 =m2[kk+3,jj]
                        alpha31 :: Float32 =m2[kk+3,jj+1]
                        for ii in (i : 1 : i+STRIDE-1) 
                            m3[ii, jj] = m3[ii, jj] + m1[ii, kk] * alpha00 + m1[ii, kk+1] * alpha10 + m1[ii, kk+2] * alpha20 + m1[ii,kk+3]*alpha30
                            m3[ii, jj+1] = m3[ii, jj+1] + m1[ii, kk] * alpha01 + m1[ii, kk+1] * alpha11 + m1[ii, kk+2]*alpha21 + m1[ii,kk+3]*alpha31 
                        end
                    end
                end
            end
        end
    end

    return 0. :: Float32
end


@debugprint "starpu_init"
starpu_init()

function multiply_with_starpu(A :: Matrix{Float32}, B :: Matrix{Float32}, C :: Matrix{Float32}, nslicesx, nslicesy)
    scale= 3
    tmin=0
    vert = StarpuDataFilter(STARPU_MATRIX_FILTER_VERTICAL_BLOCK, nslicesx)
    horiz = StarpuDataFilter(STARPU_MATRIX_FILTER_BLOCK, nslicesy)
    @starpu_block let
        hA,hB,hC = starpu_data_register(A, B, C)
        starpu_data_partition(hB, vert)
        starpu_data_partition(hA, horiz)
        starpu_data_map_filters(hC, vert, horiz)
        tmin=0
        perfmodel = StarpuPerfmodel(
            perf_type = STARPU_HISTORY_BASED,
            symbol = "history_perf"
        )
        cl = StarpuCodelet(
            cpu_func = CPU_CODELETS["matrix_mult"],
            #cuda_func = "matrix_mult",
            #opencl_func="ocl_matrix_mult",
            modes = [STARPU_R, STARPU_R, STARPU_W],
            perfmodel = perfmodel
        )

        for i in (1 : 10 )
            t=time_ns()
            @starpu_sync_tasks begin
                for taskx in (1 : nslicesx)
                    for tasky in (1 : nslicesy)
                        handles = [hA[tasky], hB[taskx], hC[taskx, tasky]]
                        task = StarpuTask(cl = cl, handles = handles)
                        starpu_task_submit(task)
                        #@starpu_async_cl matrix_mult(hA[tasky], hB[taskx], hC[taskx, tasky])
                    end
                end
            end
            t=time_ns()-t
            if (tmin==0 || tmin>t)
                tmin=t
            end
        end
    end
    return tmin
end


function approximately_equals(
    A :: Matrix{Cfloat},
    B :: Matrix{Cfloat},
    eps = 1e-2
)
    (height, width) = size(A)

    for j in (1 : width)
        for i in (1 : height)
            if (abs(A[i,j] - B[i,j]) > eps * max(abs(B[i,j]), abs(A[i,j])))
                println("A[$i,$j] : $(A[i,j]), B[$i,$j] : $(B[i,j])")
                return false
            end
        end
    end

    return true
end

function compute_times(io,start_dim, step_dim, stop_dim, nslicesx, nslicesy)
    for dim in (start_dim : step_dim : stop_dim)
        A = Array(rand(Cfloat, dim, dim))
        B = Array(rand(Cfloat, dim, dim))
        C = zeros(Float32, dim, dim)
        mt =  multiply_with_starpu(A, B, C, nslicesx, nslicesy)
        flops = (2*dim-1)*dim*dim/mt
        size=dim*dim*4*3/1024/1024
        println(io,"$size $flops")
        println("$size $flops")
    end
end


io=open(ARGS[1],"w")
compute_times(io,16*STRIDE,4*STRIDE,4096,2,2)
close(io)
@debugprint "starpu_shutdown"
starpu_shutdown()