<commit_before>#include <tiramisu/tiramisu.h>

#include "configuration.h"

#define GEMM_BATCH 10

using namespace tiramisu;

int main(int argc, char **argv)
{
    // Single LSTM block without minibatching
    tiramisu::init("lstm");

    // -------------------------------------------------------
    // Layer I
    // -------------------------------------------------------

    // Inner dimensions
    var i("i", 0, FEATURE_SIZE), j("j", 0, FEATURE_SIZE), k("k", 0, BATCH_SIZE);
    var i_merged("i_merged", 0, 4 * FEATURE_SIZE);
    var i0("i0"), i1("i1"), k0("k0"), k1("k1");
    // Outer dimensions
    var l("l", 0, NUM_LAYERS), s("s", 0, SEQ_LENGTH);
    var s0("s0", 0, SEQ_LENGTH / GEMM_BATCH), s1("s1", 0, GEMM_BATCH);
    // After skewing
    var l_s("l_s"), s_s("s_s");

    input x("x", {s, k, i}, p_float32);
    input weights("weights", {l, var("w_i", 0, 2), i_merged, j}, p_float32);
    input biases("biases", {l, i_merged}, p_float32);
    input tmp("tmp", {s, k, i_merged}, p_float32);
    weights.get_buffer()->tag_gpu_global();
    biases.get_buffer()->tag_gpu_global();
    tmp.get_buffer()->tag_gpu_global();

    buffer buf_Weights_cpu("buf_Weights_cpu", {NUM_LAYERS, 2, 4 * FEATURE_SIZE, FEATURE_SIZE}, p_float32, a_input);
    buffer buf_h("buf_h", {NUM_LAYERS + 1, SEQ_LENGTH + 1, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_temporary);
    buffer buf_biases_cpu("buf_biases_cpu", {NUM_LAYERS, 4 * FEATURE_SIZE}, p_float32, a_input);
    buffer buf_x_cpu("buf_x_cpu", {SEQ_LENGTH, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_input);
    buffer buf_y_cpu("buf_y_cpu", {SEQ_LENGTH, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_output);
    buffer buf_x("buf_x", {SEQ_LENGTH, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_temporary);
    buffer buf_y("buf_y", {SEQ_LENGTH, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_temporary);
    buffer buf_tmp_i("buf_tmp_i", {NUM_LAYERS, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_temporary);
    buffer buf_tmp_z("buf_tmp_z", {NUM_LAYERS, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_temporary);
    buffer buf_tmp_o("buf_tmp_o", {NUM_LAYERS, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_temporary);
    buffer buf_tmp_f("buf_tmp_f", {NUM_LAYERS, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_temporary);
    buffer buf_c("buf_c", {NUM_LAYERS, SEQ_LENGTH + 1, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_temporary);
    buf_h.tag_gpu_global();
    buf_x.tag_gpu_global();
    buf_y.tag_gpu_global();
    buf_tmp_i.tag_gpu_global();
    buf_tmp_z.tag_gpu_global();
    buf_tmp_o.tag_gpu_global();
    buf_tmp_f.tag_gpu_global();
    buf_c.tag_gpu_global();

    // Transpose Weights
    var w_t_i("w_i", 0, 2);  // Dummy variable
    computation weights_T({l, w_t_i, j, i_merged}, weights(l, w_t_i, i_merged, j));
    weights_T.get_buffer()->tag_gpu_global();
    // h(l, s) is the output of the block (l, s)
    // which takes h(l, s - 1) and h(l - 1, s) as inputs
    // initial hidden states are h(l, -1) and c(l, -1)
    // input x is copied to h(-1, s)
    computation h({l, s, k, i}, p_float32);
    computation c({l, s, k, i}, p_float32);
    computation h_init({l, k, i}, expr(float(0)));
    computation c_init({l, k, i}, expr(float(0)));
    computation h_copy_x({s, k, i}, x(s, k, i));
    // Multiplication from input is 2xbatched:
    computation sum1({l, s0},
        cublas_sgemm(buf_h, *weights_T.get_buffer(), *tmp.get_buffer(),
                     GEMM_BATCH * BATCH_SIZE, 4 * FEATURE_SIZE, FEATURE_SIZE,
                     1, 0,  // alpha, beta
                     0, 0, 0,  // ldABC
                     (l * (SEQ_LENGTH + 1) + s0 * GEMM_BATCH + 1) * BATCH_SIZE * FEATURE_SIZE,  //offsetA
                     (l * 2 + 1) * 4 * FEATURE_SIZE * FEATURE_SIZE,  //offsetB
                     s0 * GEMM_BATCH * BATCH_SIZE * 4 * FEATURE_SIZE,  // offsetC
                     false, false));
    computation sum2({l, s},
        cublas_sgemm(buf_h, *weights_T.get_buffer(), *tmp.get_buffer(),
                     BATCH_SIZE, 4 * FEATURE_SIZE, FEATURE_SIZE,
                     1, 1,  // alpha, beta
                     0, 0, 0,  // ldABC
                     ((l + 1) * (SEQ_LENGTH + 1) + s) * BATCH_SIZE * FEATURE_SIZE,  //offsetA
                     (l * 2) * 4 * FEATURE_SIZE * FEATURE_SIZE,  //offsetB
                     s * BATCH_SIZE * 4 * FEATURE_SIZE,  // offsetC
                     false, false));
    #define sigmoid(x) expr(float(1)) / (1 + expr(o_expo, -(x)))
    computation sig_i({l, s, k, i},      sigmoid(tmp(s, k, i + 0 * FEATURE_SIZE) + biases(l, i + 0 * FEATURE_SIZE)));
    computation tnh_z({l, s, k, i}, expr(o_tanh, tmp(s, k, i + 1 * FEATURE_SIZE) + biases(l, i + 1 * FEATURE_SIZE)));
    computation sig_o({l, s, k, i},      sigmoid(tmp(s, k, i + 2 * FEATURE_SIZE) + biases(l, i + 2 * FEATURE_SIZE)));
    computation sig_f({l, s, k, i},      sigmoid(tmp(s, k, i + 3 * FEATURE_SIZE) + biases(l, i + 3 * FEATURE_SIZE)));
    computation mul_iz({l, s, k, i}, sig_i(l, s, k, i) * tnh_z(l, s, k, i));
    computation mul_fc({l, s, k, i}, sig_f(l, s, k, i) * c(l, s - 1, k, i));
    c.set_expression(mul_iz(l, s, k, i) + mul_fc(l, s, k, i));
    computation tnh_c({l, s, k, i}, expr(o_tanh, c(l, s, k, i)));
    h.set_expression(tnh_c(l, s, k, i) * sig_o(l, s, k, i));
    computation stream_sync({l, s0}, cuda_stream_synchronize());
    // Output is the last layer
    computation y({s, k, i}, h(NUM_LAYERS - 1, s, k, i));
    // Copies
    computation copy_Weights_to_device({}, memcpy(buf_Weights_cpu, *weights.get_buffer()));
    computation copy_biases_to_device({}, memcpy(buf_biases_cpu, *biases.get_buffer()));
    computation copy_x_to_device({}, memcpy(buf_x_cpu, buf_x));
    computation copy_y_to_host({}, memcpy(buf_y, buf_y_cpu));

    // -------------------------------------------------------
    // Layer II
    // -------------------------------------------------------

    weights_T.gpu_tile(j, i_merged, 16, 16);
    block nonlinear_block({&sig_i, &tnh_z, &sig_o, &sig_f, &mul_iz, &mul_fc, &c,
                           &tnh_c, &h});<commit_msg>LSTM_lib kernel fusion<commit_after>#include <tiramisu/tiramisu.h>

#include "configuration.h"

#define GEMM_BATCH 10

using namespace tiramisu;

int main(int argc, char **argv)
{
    // Single LSTM block without minibatching
    tiramisu::init("lstm");

    // -------------------------------------------------------
    // Layer I
    // -------------------------------------------------------

    // Inner dimensions
    var i("i", 0, FEATURE_SIZE), j("j", 0, FEATURE_SIZE), k("k", 0, BATCH_SIZE);
    var i_merged("i_merged", 0, 4 * FEATURE_SIZE);
    var i0("i0"), i1("i1"), k0("k0"), k1("k1");
    // Outer dimensions
    var l("l", 0, NUM_LAYERS), s("s", 0, SEQ_LENGTH);
    var s0("s0", 0, SEQ_LENGTH / GEMM_BATCH), s1("s1", 0, GEMM_BATCH);
    // After skewing
    var l_s("l_s"), s_s("s_s");

    input x("x", {s, k, i}, p_float32);
    input weights("weights", {l, var("w_i", 0, 2), i_merged, j}, p_float32);
    input biases("biases", {l, i_merged}, p_float32);
    input tmp("tmp", {s, k, i_merged}, p_float32);
    weights.get_buffer()->tag_gpu_global();
    biases.get_buffer()->tag_gpu_global();
    tmp.get_buffer()->tag_gpu_global();

    buffer buf_Weights_cpu("buf_Weights_cpu", {NUM_LAYERS, 2, 4 * FEATURE_SIZE, FEATURE_SIZE}, p_float32, a_input);
    buffer buf_h("buf_h", {NUM_LAYERS + 1, SEQ_LENGTH + 1, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_temporary);
    buffer buf_biases_cpu("buf_biases_cpu", {NUM_LAYERS, 4 * FEATURE_SIZE}, p_float32, a_input);
    buffer buf_x_cpu("buf_x_cpu", {SEQ_LENGTH, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_input);
    buffer buf_y_cpu("buf_y_cpu", {SEQ_LENGTH, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_output);
    buffer buf_x("buf_x", {SEQ_LENGTH, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_temporary);
    buffer buf_y("buf_y", {SEQ_LENGTH, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_temporary);
    buffer buf_tmp_i("buf_tmp_i", {NUM_LAYERS, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_temporary);
    buffer buf_tmp_z("buf_tmp_z", {NUM_LAYERS, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_temporary);
    buffer buf_tmp_o("buf_tmp_o", {NUM_LAYERS, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_temporary);
    buffer buf_tmp_f("buf_tmp_f", {NUM_LAYERS, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_temporary);
    buffer buf_c("buf_c", {NUM_LAYERS, SEQ_LENGTH + 1, BATCH_SIZE, FEATURE_SIZE}, p_float32, a_temporary);
    buf_h.tag_gpu_global();
    buf_x.tag_gpu_global();
    buf_y.tag_gpu_global();
    buf_tmp_i.tag_gpu_global();
    buf_tmp_z.tag_gpu_global();
    buf_tmp_o.tag_gpu_global();
    buf_tmp_f.tag_gpu_global();
    buf_c.tag_gpu_global();

    // Transpose Weights
    var w_t_i("w_i", 0, 2);  // Dummy variable
    computation weights_T({l, w_t_i, j, i_merged}, weights(l, w_t_i, i_merged, j));
    weights_T.get_buffer()->tag_gpu_global();
    // h(l, s) is the output of the block (l, s)
    // which takes h(l, s - 1) and h(l - 1, s) as inputs
    // initial hidden states are h(l, -1) and c(l, -1)
    // input x is copied to h(-1, s)
    computation h({l, s, k, i}, p_float32);
    computation c({l, s, k, i}, p_float32);
    computation h_init({l, k, i}, expr(float(0)));
    computation c_init({l, k, i}, expr(float(0)));
    computation h_copy_x({s, k, i}, x(s, k, i));
    // Multiplication from input is batched:
    computation sum1({l, s0},
        cublas_sgemm(buf_h, *weights_T.get_buffer(), *tmp.get_buffer(),
                     GEMM_BATCH * BATCH_SIZE, 4 * FEATURE_SIZE, FEATURE_SIZE,
                     1, 0,  // alpha, beta
                     0, 0, 0,  // ldABC
                     (l * (SEQ_LENGTH + 1) + s0 * GEMM_BATCH + 1) * BATCH_SIZE * FEATURE_SIZE,  //offsetA
                     (l * 2 + 1) * 4 * FEATURE_SIZE * FEATURE_SIZE,  //offsetB
                     s0 * GEMM_BATCH * BATCH_SIZE * 4 * FEATURE_SIZE,  // offsetC
                     false, false));
    computation sum2({l, s},
        cublas_sgemm(buf_h, *weights_T.get_buffer(), *tmp.get_buffer(),
                     BATCH_SIZE, 4 * FEATURE_SIZE, FEATURE_SIZE,
                     1, 1,  // alpha, beta
                     0, 0, 0,  // ldABC
                     ((l + 1) * (SEQ_LENGTH + 1) + s) * BATCH_SIZE * FEATURE_SIZE,  //offsetA
                     (l * 2) * 4 * FEATURE_SIZE * FEATURE_SIZE,  //offsetB
                     s * BATCH_SIZE * 4 * FEATURE_SIZE,  // offsetC
                     false, false));
    #define sigmoid(x) expr(float(1)) / (1 + expr(o_expo, -(x)))
    computation sig_i({l, s, k, i},      sigmoid(tmp(s, k, i + 0 * FEATURE_SIZE) + biases(l, i + 0 * FEATURE_SIZE)));
    computation tnh_z({l, s, k, i}, expr(o_tanh, tmp(s, k, i + 1 * FEATURE_SIZE) + biases(l, i + 1 * FEATURE_SIZE)));
    computation sig_o({l, s, k, i},      sigmoid(tmp(s, k, i + 2 * FEATURE_SIZE) + biases(l, i + 2 * FEATURE_SIZE)));
    computation sig_f({l, s, k, i},      sigmoid(tmp(s, k, i + 3 * FEATURE_SIZE) + biases(l, i + 3 * FEATURE_SIZE)));
    computation mul_iz({l, s, k, i}, sig_i(l, s, k, i) * tnh_z(l, s, k, i));
    computation mul_fc({l, s, k, i}, sig_f(l, s, k, i) * c(l, s - 1, k, i));
    c.set_expression(mul_iz(l, s, k, i) + mul_fc(l, s, k, i));
    computation tnh_c({l, s, k, i}, expr(o_tanh, c(l, s, k, i)));
    h.set_expression(tnh_c(l, s, k, i) * sig_o(l, s, k, i));
    computation stream_sync({l, s0}, cuda_stream_synchronize());
    // Output is the last layer
    computation y({s, k, i}, h(NUM_LAYERS - 1, s, k, i));
    // Copies
    computation copy_Weights_to_device({}, memcpy(buf_Weights_cpu, *weights.get_buffer()));
    computation copy_biases_to_device({}, memcpy(buf_biases_cpu, *biases.get_buffer()));
    computation copy_x_to_device({}, memcpy(buf_x_cpu, buf_x));
    computation copy_y_to_host({}, memcpy(buf_y, buf_y_cpu));

    // -------------------------------------------------------
    // Layer II
    // -------------------------------------------------------

    // Fuse kernels by moving gpu iterators out
    weights_T.interchange(w_t_i, j);
    weights_T.interchange(w_t_i, i_merged);
    weights_T.interchange(l, j);
    weights_T.interchange(l, i_merged);
    h_init.interchange(l, k);
    h_init.interchange(l, i);
    c_init.interchange(l, k);
    c_init.interchange(l, i);
    h_copy_x.interchange(s, k);
    h_copy_x.interchange(s, i);
    y.interchange(s, k);
    y.interchange(s, i);

    weights_T.gpu_tile(j, i_merged, 16, 16);
    block nonlinear_block({&sig_i, &tnh_z, &sig_o, &sig_f, &mul_iz, &mul_fc, &c,
                           &tnh_c, &h});
{'fraction_non_alphanumeric': 0.19148936170212766, 'fraction_numerical': 0.019436890097584387, 'mean_word_length': 2.694739952718676, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1684150', 'n_tokens_mistral': 5043, 'n_tokens_neox': 4834, 'n_words': 1098}