kutacc_af2_triangle_multiplication_last

Last step of triangle_multiplication, where an output tensor of out × gate is generated according to the step out_linear.

Interface Definition

void kutacc_af2_triangle_multiplication_last(kutacc_tensor_h out, kutacc_tensor_h gate, int64_t n_res, int64_t n_res_gather, int64_t c_o);

Parameters

Table 1 Input parameter definition

Parameter	Type	Description	Input/Output
out	kutacc_tensor_h	Output data	Output
gate	kutacc_tensor_h	Gating tensor	Input
n_res	int64_t	Number of residues	Input
n_res_gather	int64_t	Number of residues	Input
c_o	int64_t	Output feature dimension	Input

Constraints on integer parameters of triangle_multiplication:

n_res, n_res_gather, c_o > 0

n_res * n_res_gather < INT64_MAX

c_i = c_o

In single-process mode, n_res must be equal to n_res_gather. In multi-process mode, this condition does not need to be met.

The table below describes operator shape constraints when a test case is built or when KPEX is called.

Table 2 Shape constraints on kpex triangle_multiplication parameters

Tensor	Shape	Description
act	[n_res, n_res_gather, c_z]	KPEX layer input, where the size of the last dimension of act is equal to c_o and c_i.
mask	[n_res, n_res_gather]	Mask input of the KPEX layer
input_ln_w	[c_o]	Weight of act generated after LayerNorm operation
input_ln_b	[c_o]	Bias of act generated after LayerNorm operation
left_proj_w	[c_i, c_o]	See the definition of proj_w in "Table 2 kutacc_af2_tm_proj_weights_t structure definition" in section "8.1.2.2.2.1.1.2".
left_proj_b	[c_i]	See the definition of proj_b in "Table 2 kutacc_af2_tm_proj_weights_t structure definition" in section "8.1.2.2.2.1.1.2".
right_proj_w	[c_i, c_o]	See the definition of proj_w in "Table 2 kutacc_af2_tm_proj_weights_t structure definition" in section "8.1.2.2.2.1.1.2".
right_proj_b	[c_i]	See the definition of proj_b in "Table 2 kutacc_af2_tm_proj_weights_t structure definition" in section "8.1.2.2.2.1.1.2".
left_gate_w	[c_i, c_o]	See the definition of gate_w in "Table 2 kutacc_af2_tm_proj_weights_t structure definition" in section "8.1.2.2.2.1.1.2".
left_gate_b	[c_i]	See the definition of gate_b in "Table 2 kutacc_af2_tm_proj_weights_t structure definition" in section "8.1.2.2.2.1.1.2".
right_gate_w	[c_i, c_o]	See the definition of gate_w in "Table 2 kutacc_af2_tm_proj_weights_t structure definition" in section "8.1.2.2.2.1.1.2".
right_gate_b	[c_i]	See the definition of gate_b in "Table 2 kutacc_af2_tm_proj_weights_t structure definition" in section "8.1.2.2.2.1.1.2".
gating_w	[c_i, c_o]	See the definition of gating_w in "Table 3 kutacc_af2_tm_linear_weights_t structure definition" in the previous section.
gating_b	[c_i]	See the definition of gating_b in "Table 3 kutacc_af2_tm_linear_weights_t structure definition" in the previous section.
center_ln_w	[c_i]	Weight of center_act during LayerNorm normalization process.
center_ln_b	[c_i]	Bias of center_act during LayerNorm normalization process.
output_proj_w	[c_o, c_i]	See the definition of output_proj_w in "Table 3 kutacc_af2_tm_linear_weights_t structure definition" in the previous section.
output_proj_b	[c_o]	See the definition of output_proj_b in "Table 3 kutacc_af2_tm_linear_weights_t structure definition" in the previous section.

Examples

C++ interface:

// test_triangle_multiplication.h
#define KPEX_TPP_ALPHAFOLD_TEST_TMP_H
#include <ATen/core/Tensor.h>
#include <ATen/ops/empty.h>
#include <ATen/ops/ones.h>
#include <ATen/ops/zeros.h>
#include <ATen/ops/full.h>
#include <ATen/native/cpu/utils.h>
#include <c10/core/ScalarType.h>
namespace alphafold {
at::Tensor test_triangle_multiplication(int64_t n_res, int64_t n_res_gather, int64_t c_o, int64_t c_i);
}
#endif

// bind.h
#include <torch/extension.h>
#include "test_triangle_multiplication.h"
namespace alphafold {
inline void bind(pybind11::module &m)
{
    auto submodule = m.def_submodule("alphafold");
    submodule.def("test_triangle_multiplication", &test_triangle_multiplication, 
        py::arg("n_res"), py::arg("n_res_gather"), py::arg("c_o"), py::arg("c_i"));
}
}

// test.py
import torch
from torch import nn
import numpy as np
import torch.distributed as dist
import kpex._C as kernel
import kpex
import os
def test_triangle_multiplication(n_res, n_res_gather, c_o, c_i):
    out = kernel.alphafold.test_triangle_multiplication(n_res, n_res_gather, c_o, c_i)
    return out

// test_triangle_multiplication.cpp
#include "test_triangle_multiplication.h"
#include "triangle_multiplication.h"
#include "kutacc.h"
#include "utils/memory.h"
#include "utils/layernorm.h"
#include "utils/TensorWrapper.h"
namespace alphafold {
at::Tensor test_triangle_multiplication(int64_t n_res, int64_t n_res_gather, int64_t c_o, int64_t c_i)
{
    float a = 0.5f;
    float b = 0.2f;
    float c = 0.4f;
    float d = 1.25f;
    auto bf16_opt = at::TensorOptions().device(kpex::device()).dtype(at::kBFloat16);
    auto float_opt = at::TensorOptions().device(kpex::device()).dtype(at::kFloat);
    at::Tensor input_ln_w = at::full({c_o}, a, float_opt);
    at::Tensor input_ln_b = at::full({c_o}, b, float_opt);
    at::Tensor left_proj_w = linear_weight_prepack(at::ones({c_i, c_o}, bf16_opt));
    at::Tensor left_proj_b = at::zeros({c_i}, float_opt);
    at::Tensor right_proj_w = linear_weight_prepack(at::full({c_i, c_o}, c, bf16_opt));
    at::Tensor right_proj_b = at::ones({c_i}, float_opt);
    at::Tensor left_gate_w = linear_weight_prepack(at::ones({c_i, c_o}, bf16_opt));
    at::Tensor left_gate_b = at::ones({c_i}, float_opt);
    at::Tensor right_gate_w = linear_weight_prepack(at::zeros({c_i, c_o}, bf16_opt));
    at::Tensor right_gate_b = at::zeros({c_i}, float_opt);
    at::Tensor gating_w = linear_weight_prepack(at::full({c_i, c_o}, d, bf16_opt));
    at::Tensor gating_b = at::zeros({c_i}, float_opt);
    at::Tensor center_ln_w = at::ones({c_i}, float_opt);
    at::Tensor center_ln_b = at::zeros({c_i}, float_opt);
    at::Tensor output_proj_w = linear_weight_prepack(at::zeros({c_o, c_i}, bf16_opt));
    at::Tensor output_proj_b =  at::zeros({c_o}, float_opt);
    float e = 1.2f;
    at::Tensor act = at::full({n_res, n_res_gather, c_o}, e, bf16_opt);
    at::Tensor mask = at::ones({n_res, n_res_gather}, bf16_opt);
    at::Tensor out = at::empty(act.sizes(), act.options());
    at::Tensor input_act = layernorm(act, input_ln_w, input_ln_b);
    at::Tensor center_act;
    at::Tensor left_proj_act = input_act.new_empty({c_i, n_res, n_res_gather});
    at::Tensor right_proj_act = input_act.new_empty({c_i, n_res, n_res_gather});
    at::Tensor gate = act.new_empty({n_res, n_res_gather, c_o});
    auto left_proj_act_tw = convert_to_tensor_wrapper(left_proj_act);
    auto right_proj_act_tw = convert_to_tensor_wrapper(right_proj_act);
    auto gate_tw = convert_to_tensor_wrapper(gate);
    auto input_act_tw = convert_to_tensor_wrapper(input_act);
    auto mask_tw = convert_to_tensor_wrapper(mask);
    auto left_proj_w_tw = convert_to_tensor_wrapper(left_proj_w);
    auto left_proj_b_tw = convert_to_tensor_wrapper(left_proj_b);
    auto left_gate_w_tw = convert_to_tensor_wrapper(left_gate_w);
    auto left_gate_b_tw = convert_to_tensor_wrapper(left_gate_b);
    auto right_proj_w_tw = convert_to_tensor_wrapper(right_proj_w);
    auto right_proj_b_tw = convert_to_tensor_wrapper(right_proj_b);
    auto right_gate_w_tw= convert_to_tensor_wrapper(right_gate_w);
    auto right_gate_b_tw= convert_to_tensor_wrapper(right_gate_b);
    at::Tensor gate_left = input_act.new_empty({c_i, n_res, n_res_gather});
    auto gate_left_tw = convert_to_tensor_wrapper(gate_left);
    kutacc_af2_tm_act_inputs_t_wrapper *left_acts_ptr = new kutacc_af2_tm_act_inputs_t_wrapper(left_proj_act_tw, 
        input_act_tw, gate_left_tw, n_res, n_res_gather);
    kutacc_af2_tm_proj_weights_t_wrapper *left_weights_ptr = new kutacc_af2_tm_proj_weights_t_wrapper(left_proj_w_tw,
        left_proj_b_tw, left_gate_w_tw, left_gate_b_tw, c_o, c_i);
    kutacc_af2_triangle_multiplication_calc_proj(left_acts_ptr, mask_tw.get_tensor(), left_weights_ptr, false);
    
    at::Tensor gate_right = input_act.new_empty({c_i, n_res, n_res_gather});
    auto gate_right_tw = convert_to_tensor_wrapper(gate_right);
    kutacc_af2_tm_act_inputs_t_wrapper *right_acts_ptr = new kutacc_af2_tm_act_inputs_t_wrapper(right_proj_act_tw, 
        input_act_tw, gate_right_tw, n_res, n_res_gather);
    kutacc_af2_tm_proj_weights_t_wrapper *right_weights_ptr = new kutacc_af2_tm_proj_weights_t_wrapper(right_proj_w_tw, 
        right_proj_b_tw, right_gate_w_tw, right_gate_b_tw, c_o, c_i);
    kutacc_af2_triangle_multiplication_calc_proj(right_acts_ptr, mask_tw.get_tensor(), right_weights_ptr, false);
    center_act = act.new_empty({left_proj_act.sizes()[0], n_res_gather, n_res_gather});
    auto center_act_tw = convert_to_tensor_wrapper(center_act);
    kutacc_af2_triangle_multiplication_equation(center_act_tw.get_tensor(), left_proj_act_tw.get_tensor(), 
        right_proj_act_tw.get_tensor(), n_res_gather, true);
    center_act = center_act.permute({1, 2, 0}).contiguous();
    center_act = layernorm(center_act, center_ln_w, center_ln_b);
    auto center_act_new_tw = convert_to_tensor_wrapper(center_act);
    auto out_tw = convert_to_tensor_wrapper(out);
    auto gating_w_tw = convert_to_tensor_wrapper(gating_w);
    auto gating_b_tw = convert_to_tensor_wrapper(gating_b);
    auto output_proj_w_tw = convert_to_tensor_wrapper(output_proj_w);
    auto output_proj_b_tw = convert_to_tensor_wrapper(output_proj_b);
    kutacc_af2_tm_linear_weights_t_wrapper *linear_weights_ptr = new kutacc_af2_tm_linear_weights_t_wrapper(gating_w_tw,
        gating_b_tw, output_proj_w_tw, output_proj_b_tw, c_o, c_i);
    kutacc_af2_triangle_multiplication_gate_and_out_linear(gate_tw.get_tensor(), out_tw.get_tensor(), left_acts_ptr, 
        center_act_new_tw.get_tensor(), linear_weights_ptr, true);
    kutacc_af2_triangle_multiplication_last(out_tw.get_tensor(), gate_tw.get_tensor(), n_res, n_res_gather, c_o);
    delete left_acts_ptr;
    delete left_weights_ptr;
    delete right_acts_ptr;
    delete right_weights_ptr;
    delete linear_weights_ptr;
    return out;
}
}

Parent topic: AlphaFold2