kutacc_af2_outer_product_mean_chunk

Calculate the result by block in outer_product_mean.

Interface Definition

kutacc_export void kutacc_af2_outer_product_mean_chunk(kutacc_af2_opm_act_inputs_t *opm_acts_ptr, kutacc_af2_opm_mask_inputs_t *opm_masks_ptr, kutacc_af2_opm_weights_t *opm_weights_ptr,kutacc_tensor_h out, int64_t left_block_size, int64_t right_block_size);

Parameters

Table 1 Input parameter definition

Parameter	Type	Description	Input/Output
opm_acts_ptr	kutacc_af2_opm_act_inputs_t *	Pointer of the kutacc_af2_opm_act_inputs_t type. For details about the data structure, see Table 2.	Input
opm_masks_ptr	kutacc_af2_opm_mask_inputs_t *	Pointer of the kutacc_af2_opm_mask_inputs_t type. For details about the data structure, see Table 3.	Input
opm_weights_ptr	kutacc_af2_opm_weights_t *	Pointer of the kutacc_af2_opm_weights_t type. For details about the data structure, see Table 4.	Input
out	kutacc_tensor_h	Output data	Output
left_block_size	int64_t	Left block size	Input
right_block_size	int64_t	Right block size	Input

Table 2 kutacc_af2_opm_act_inputs_t structure definition

Parameter	Type	Description	Input/Output
n_seq	int64_t	Number of sequences	Input
n_res	int64_t	Number of residues	Input
input_act	kutacc_tensor_h	Input activation tensor	Input
left_proj	kutacc_tensor_h	Left projection	Input
right_proj	kutacc_tensor_h	Right projection	Input
left_proj_	kutacc_tensor_h	Left projection after mask processing	Input
right_proj_	kutacc_tensor_h	Right projection after mask processing	Input

Table 3 kutacc_af2_opm_mask_inputs_t structure definition

Parameter	Type	Description	Input/Output
n_res_gather	int64_t	Number of residues after aggregation	Input
mask_bias	int64_t	Mask tensor address bias	Input
mask	kutacc_tensor_h	Mask tensor	Input
norm	kutacc_tensor_h	Normalization factor tensor	Input

Table 4 kutacc_af2_opm_weights_t structure definition

Parameter	Type	Description	Input/Output
c_m	int64_t	Input feature dimension	Input
c_i	int64_t	Feature dimension after projection	Input
c_z	int64_t	Output feature dimension	Input
left_proj_w	kutacc_tensor_h	Left projection weight	Input
left_proj_b	kutacc_tensor_h	Left projection bias	Input
right_proj_w	kutacc_tensor_h	Right projection weight	Input
right_proj_b	kutacc_tensor_h	Right projection bias	Input
outer_w	kutacc_tensor_h	Output weight	Input
outer_b	kutacc_tensor_h	Output bias	Input

Constraints on integer parameters of outer_product_mean:

n_res, n_res_gather, c_i, c_z, n_res, n_res_gather, left_block_size, right_block_size > 0;

n_seq * n_res < INT64_MAX

left_block_size * right_block_size * c_i * c_i < INT64_MAX

left_block_size* right_block_size * c_z < INT64_MAX

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

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

Table 5 Shape constraints on KPEX outer_product_means

tensor/param	shape/value	Description
input_ln_w	[c_m]	Weight required for input_act generated through LayerNorm
input_ln_b	[c_m]	Bias required for input_act generated through LayerNorm
left_proj_w	[c_i, c_m]	See left_proj_w in Table 4.
left_proj_b	[c_i]	See left_proj_b in Table 4.
right_proj_w	[c_i, c_m]	See right_proj_w in Table 4.
right_proj_b	[c_i]	See right_proj_b in Table 4.
output_w	[c_z, c_i, c_i]	See outer_w in Table 4.
output_b	[c_z]	See outer_b in Table 4.
act	[n_seq, n_res, c_m]	KPEX input, which then functions as input_act after LayerNorm operation.
mask	[n_seq, n_res_gather]	See mask in Table 3.
left_block_size	greater than 0 or None	See left_block_size in Table 1.
right_block_size	greater than 0 or None	See right_block_size in Table 2.

Examples

C++ interface:

// test_outer_product_mean.h
#ifndef KPEX_TPP_ALPHAFOLD_TEST_OPM_H
#define KPEX_TPP_ALPHAFOLD_TEST_OPM_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_outer_product_mean(int64_t c_i, int64_t c_m, int64_t c_z, int64_t n_seq, int64_t n_res, int64_t n_res_gather);
}
#endif

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

// test.py
import copy
import time
import types
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_outer_product_mean.cpp
#include "test_outer_product_mean.h"
#include "kutacc.h"
#include "outer_product_mean.h"
#include "utils/memory.h"
#include "utils/layernorm.h"
#include <utils/TensorWrapper.h>
namespace alphafold {
at::Tensor test_outer_product_mean(int64_t c_i, int64_t c_m, int64_t c_z, int64_t n_seq, int64_t n_res, int64_t n_res_gather)
{
    float a = 0.2f;
    float b = 0.5f;
    float c = 1.5f;
    float d = 2.0f;
    at::Tensor act = at::full({n_seq, n_res, c_m}, d, at::TensorOptions().device(kpex::device()).dtype(at::kBFloat16));
    at::Tensor mask = at::ones({n_res_gather, n_seq}, at::TensorOptions().device(kpex::device()).dtype(at::kBFloat16));
    at::Tensor left_proj = act.new_empty({c_i, n_res, n_seq});
    at::Tensor right_proj = act.new_empty({c_i, n_res, n_seq});
    at::Tensor left_proj_ = act.new_empty({n_res, c_i, n_seq});
    at::Tensor right_proj_ = act.new_empty({n_res, c_i, n_seq});
    at::Tensor norm = mask.new_empty({n_res, n_res_gather});
    int64_t mask_bias = 0;
    at::Tensor input_ln_w = at::full({c_m}, a, at::TensorOptions().device(kpex::device()).dtype(at::kFloat));
    at::Tensor input_ln_b = at::full({c_m}, b, at::TensorOptions().device(kpex::device()).dtype(at::kFloat));
    at::Tensor left_proj_w = linear_weight_prepack(at::full({c_i, c_m}, c, at::TensorOptions().device(kpex::device()).dtype(at::kBFloat16)));
    at::Tensor left_proj_b = at::zeros({c_i}, at::TensorOptions().device(kpex::device()).dtype(at::kFloat));
    at::Tensor right_proj_w = linear_weight_prepack(at::ones({c_i, c_m}, at::TensorOptions().device(kpex::device()).dtype(at::kBFloat16)));
    at::Tensor right_proj_b = at::zeros({c_i}, at::TensorOptions().device(kpex::device()).dtype(at::kFloat));
    at::Tensor output_w = linear_weight_prepack(at::ones({c_z, c_i * c_i}, at::TensorOptions().device(kpex::device()).dtype(at::kBFloat16)));
    at::Tensor output_b = at::zeros({c_z}, at::TensorOptions().device(kpex::device()).dtype(at::kFloat));
    at::Tensor input_act = layernorm(act.transpose(0, 1), input_ln_w, input_ln_b);
    at::Tensor out = act.new_empty({n_res, n_res_gather, c_z});
    kutacc::TensorWrapper input_act_tw = convert_to_tensor_wrapper(input_act);
    kutacc::TensorWrapper mask_tw = convert_to_tensor_wrapper(mask);
    kutacc::TensorWrapper left_proj_w_tw = convert_to_tensor_wrapper(left_proj_w);
    kutacc::TensorWrapper left_proj_b_tw = convert_to_tensor_wrapper(left_proj_b);
    kutacc::TensorWrapper right_proj_w_tw = convert_to_tensor_wrapper(right_proj_w);
    kutacc::TensorWrapper right_proj_b_tw = convert_to_tensor_wrapper(right_proj_b);
    kutacc::TensorWrapper left_proj_tw = convert_to_tensor_wrapper(left_proj);
    kutacc::TensorWrapper right_proj_tw = convert_to_tensor_wrapper(right_proj);
    kutacc::TensorWrapper left_proj_tw_ = convert_to_tensor_wrapper(left_proj_);
    kutacc::TensorWrapper right_proj_tw_ = convert_to_tensor_wrapper(right_proj_);
    kutacc::TensorWrapper norm_tw = convert_to_tensor_wrapper(norm);
    kutacc::TensorWrapper output_w_tw = convert_to_tensor_wrapper(output_w);
    kutacc::TensorWrapper output_b_tw = convert_to_tensor_wrapper(output_b);
    kutacc::TensorWrapper out_tw = convert_to_tensor_wrapper(out);
    int64_t left_block_size = 1024;
    int64_t right_block_size = 1024;
    kutacc_af2_opm_weights_t_wrapper *opm_weights_ptr = new kutacc_af2_opm_weights_t_wrapper(left_proj_w_tw, left_proj_b_tw, right_proj_w_tw, right_proj_b_tw,
        output_w_tw, output_b_tw, c_m, c_i, c_z);
    kutacc_af2_opm_act_inputs_t_wrapper *opm_inputs_ptr = new kutacc_af2_opm_act_inputs_t_wrapper(input_act_tw, left_proj_tw, right_proj_tw, left_proj_tw_,
        right_proj_tw_, n_seq, n_res);
    kutacc_af2_opm_mask_inputs_t_wrapper *opm_mask_ptr = new kutacc_af2_opm_mask_inputs_t_wrapper(mask_tw, norm_tw, n_res_gather, mask_bias);
    kutacc_af2_outer_product_mean_calc_left_and_right_mul(opm_inputs_ptr, opm_mask_ptr, opm_weights_ptr);
    kutacc_af2_outer_product_mean_chunk(opm_inputs_ptr, opm_mask_ptr, opm_weights_ptr, out_tw.get_tensor(), left_block_size, right_block_size);
    delete opm_inputs_ptr;
    delete opm_weights_ptr;
    delete opm_mask_ptr;
    return out;
}
}

Parent topic: AlphaFold2