matmul.impala

struct Mat {
    rows  : fn() -> i32,
    cols  : fn() -> i32,
    read  : fn(i32, i32) -> f32,
    write : fn(i32, i32, f32) -> (),
}

fn @get_mat(buf: Buffer, rows: i32, cols: i32) = Mat {
    rows  = @|| rows,
    cols  = @|| cols,
    read  = @|x, y|    { bitcast[&[f32],  ](buf.data)(y * cols + x) },
    write = @|x, y, v| { bitcast[&mut[f32]](buf.data)(y * cols + x) = v },
};


fn matmul_golden(C: Mat, A: Mat, B: Mat) -> () {
    for i in range(0, C.rows()) {
        for k in range(0, C.cols()) {
            let mut sum = 0:f32;
            for j in range(0, A.cols()) {
                sum += A.read(j, i) * B.read(k, j);
            }
            C.write(k, i, sum);
        }
    }
}


fn matmul_hand_tiled(C: Mat, A: Mat, B: Mat) -> () {
    // initialize C with zeros
    for y in range(0, C.rows()) {
        for x in range(0, C.cols()) {
            C.write(x, y, 0:f32);
        }
    }

    // tiling parameter
    let bs = 256;

    for i0 in range_step(0, C.rows(), bs) {
        for k0 in range_step(0, C.cols(), bs) {
            for j0 in range_step(0, A.cols(), bs) {

                for i in range(i0, min(i0 + bs, C.rows())) {
                    for k in range(k0, min(k0 + bs, C.cols())) {
                        let mut sum = C.read(k, i);
                        for j in range(j0, min(j0 + bs, A.cols())) {
                            sum += A.read(j, i) * B.read(k, j);
                        }
                        C.write(k, i, sum);
                    }
                }

            }
        }
    }
}


fn matmul_auto_tiled(C: Mat, A: Mat, B: Mat) -> () {
    fn tiled_matmul() -> () {
        fn @get_step(lvl: i32) -> i32 { // size -> 256 -> 64 -> 1
            if lvl == 0 { 256 } else if lvl == 1 { 64 } else { 1 }
        }

        fn @(?cur_lvl) tile(cur_lvl: i32, il: i32, iu: i32, kl: i32, ku: i32, jl: i32, ju: i32) -> () {
            let step = get_step(cur_lvl);
            pe_info[i32]("step size", step);

            if step == 1 {
                for i in range(il, min(iu, C.rows())) {
                    for k in range(kl, min(ku, C.cols())) {
                        let mut sum = C.read(k, i);
                        for j in range(jl, min(ju, A.cols())) {
                            sum += A.read(j, i) * B.read(k, j);
                        }
                        C.write(k, i, sum);
                    }
                }
            } else {
                for i in range_step(il, iu, step) {
                    for k in range_step(kl, ku, step) {
                        for j in range_step(jl, ju, step) {
                            tile(cur_lvl + 1, i, i + step, k, k + step, j, j + step);
                        }
                    }
                }
            }
        }

        tile(0, 0, C.rows(), 0, C.cols(), 0, A.cols());
    }

    // initialize C with zeros
    for y in range(0, C.rows()) {
        for x in range(0, C.cols()) {
            C.write(x, y, 0:f32);
        }
    }

    tiled_matmul()
}


#[export]
fn main() -> i32 {
    let l = 512;
    let m = 1024;
    let n = 2048;
    let A_buf = alloc_cpu((l * m) as i64 * sizeof[f32]()); // mxl
    let B_buf = alloc_cpu((m * n) as i64 * sizeof[f32]()); // nxm
    let C_buf = alloc_cpu((l * n) as i64 * sizeof[f32]()); // nxl
    let R_buf = alloc_cpu((l * n) as i64 * sizeof[f32]()); // nxl
    let A = get_mat(A_buf, l, m); // mxl
    let B = get_mat(B_buf, m, n); // nxm
    let C = get_mat(C_buf, l, n); // nxl
    let R = get_mat(R_buf, l, n); // nxl
    init_rand(bitcast[&mut[f32]](A_buf.data), A.cols(), A.rows());
    init_rand(bitcast[&mut[f32]](B_buf.data), B.cols(), B.rows());

    let nops = 2 as i64 * l as i64 * m as i64 * n as i64;
    fn print_gflops(time: i64) -> () {
        print_string("Timing: ");
        print_f64(time as f64 / 1000:f64);
        print_string(" ms\n");
        let gflops = (1.0e-9 * nops as f64) / (time as f64 * 1.0e-6);
        print_f64(gflops);
        print_string(" GFLOPS\n")
    }

    fn bench_fun(fun: fn() -> ()) = @|| {
        let start = get_micro_time();
        fun();
        let end = get_micro_time();
        print_gflops(end - start);
    };

    for bench_fun() {
        matmul_golden(R, A, B);
    }

    for bench_fun() {
        matmul_hand_tiled(C, A, B);
    }

    for bench_fun() {
        matmul_auto_tiled(C, A, B);
    }


    let mut passed = 0;
    for y in range(0, C.rows()) {
        for x in range(0, C.cols()) {
            if C.read(x, y) != R.read(x, y) {
                passed++;
            }
        }
    }
    if passed == 0 {
        print_string("Test PASSED!\n");
    } else {
        print_string("Test FAILED!\n");
    }

    release(A_buf);
    release(B_buf);
    release(C_buf);
    release(R_buf);

    passed
}