kutacc_af2_transition

Input an input tensor act, normalize it using LayerNorm, and then apply a linear projection for feature expansion. Next, a ReLU activation function is applied for non-linear processing, followed by another linear projection for feature compression before returning the final output.

Interface Definition

kutacc_export void kutacc_af2_transition(kutacc_af2_trans_act_inputs_t *trans_inputs_ptr, kutacc_af2_trans_weights_t * trans_weights_ptr, kutacc_tensor_h out);

Parameters

Table 1 Parameters

Parameter	Type	Description	Input/Output
trans_inputs_ptr	kutacc_af2_trans_act_inputs_t *	Pointer of the kutacc_af2_trans_act_inputs_t type. For details, see Table 2.	Input
trans_weights_ptr	kutacc_af2_trans_weights_t *	Pointer of the kutacc_af2_trans_weights_t type. For details, see Table 3.	Input
out	kutacc_tensor_h	Output data	Output

Table 2 kutacc_af2_trans_act_inputs_t structure definition

Parameter	Type	Description	Input/Output
batch	int64_t	Size of the first dimension of act	Input
n_res	int64_t	Size of the second dimension of act	Input
input_act	kutacc_tensor_h	Input matrix obtained after act is normalized by LayerNorm	Input
intermediate_act	kutacc_tensor_h	Intermediate variable generated during the first linear transformation, used as the input for the second linear transformation	Input

Table 3 kutacc_af2_trans_weights_t structure definition

Parameter	Type	Description	Input/Output
c_o	int64_t	Size of the first dimension of input_ln_w (see input_ln_w in Table 4)	Input
c_i	int64_t	Size of the first dimension of linear1_w	Input
linear1_w	kutacc_tensor_h	Weight of the first linear transformation	Input
linear1_b	kutacc_tensor_h	Bias of the first linear transformation	Input
linear2_w	kutacc_tensor_h	Weight of the second linear transformation	Input
linear2_b	kutacc_tensor_h	Bias of the second linear transformation	Input

Constraints on integer parameters of transition:

batch, n_res, c_o, c_i > 0

batch, n_res, c_o, c_i, batch * n_res < INT64_MAX

The table below lists the parameter shape constraints on parameter construction at the KPEX layer (the full name is provided in parentheses).

Table 4 Input tensor shapes and descriptions of KPEX transition

Parameter	Shape	Description
input_ln_w(input_layer_norm_weight)	[c_o]	Weight during LayerNorm normalization of act
input_ln_b(input_layer_norm_bias)	[c_o]	Bias during LayerNorm normalization of act
linear_1_w(transition1_weight)	[c_i, c_o]	See linear1_w in Table 3.
linear_1_b(transition1_bias)	[c_i]	See linear1_b in Table 3.
linear_2_w(transition2_weight)	[c_o, c_i]	See linear2_w in Table 3.
linear_2_b(transition2_bias)	[c_o]	See linear2_b in Table 3.
act	[batch, n_res, c_o]	Input tensor.

Examples

C++ interface:

// test_transition.h
#ifndef KPEX_TPP_ALPHAFOLD_TEST_TRANSITION_H
#define KPEX_TPP_ALPHAFOLD_TEST_TRANSITION_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_transition(int64_t batch, int64_t n_res, int64_t c_o, int64_t c_i);
}
#endif

// test_transition.cpp
#include "kutacc.h"
#include "test_transition.h"
#include "transition.h"
#include <utils/memory.h>
#include "utils/TensorWrapper.h"
#include "utils/layernorm.h"
namespace alphafold {
at::Tensor test_transition(int64_t batch, int64_t n_res, int64_t c_o, int64_t c_i)
{
    float a = 0.2f;
    at::Tensor act = at::full({batch, n_res, c_o}, a, at::TensorOptions().device(kpex::device()).dtype(at::kBFloat16));
    at::Tensor out = act.new_empty(act.sizes());
    float b = 0.1f;
    float c = 0.5f;
    at::Tensor input_ln_w = at::ones({c_o}, at::TensorOptions().device(kpex::device()).dtype(at::kFloat));
    at::Tensor input_ln_b = at::zeros({c_o}, input_ln_w.options());
    at::Tensor linear1_w = at::full({c_i, c_o}, b, act.options());
    at::Tensor linear1_b = at::full({c_i}, c, input_ln_w.options());
    at::Tensor linear1_w_new = linear_weight_prepack(linear1_w);
    at::Tensor linear2_w = at::full({c_o, c_i}, b, act.options());
    at::Tensor linear2_w_new = linear_weight_prepack(linear2_w);
    at::Tensor linear_2_b = at::zeros({c_o}, input_ln_w.options());
    at::Tensor input_act = layernorm(act, input_ln_w, input_ln_b);
    at::Tensor intermediate_act = act.new_empty({batch * n_res, c_i});
    kutacc::TensorWrapper input_act_tw = convert_to_tensor_wrapper(input_act);
    kutacc::TensorWrapper linear1_w_tw = convert_to_tensor_wrapper(linear1_w_new);
    kutacc::TensorWrapper linear1_b_tw = convert_to_tensor_wrapper(linear1_b);
    kutacc::TensorWrapper linear2_w_tw = convert_to_tensor_wrapper(linear2_w_new);
    kutacc::TensorWrapper linear2_b_tw = convert_to_tensor_wrapper(linear2_b);
    kutacc::TensorWrapper intermediate_act_tw = convert_to_tensor_wrapper(intermediate_act);
    kutacc::TensorWrapper out_tw = convert_to_tensor_wrapper(out);
    
    kutacc_af2_trans_weights_t_wrapper *trans_weights_ptr = new kutacc_af2_trans_weights_t_wrapper(linear1_w_tw, linear1_b_tw, linear2_w_tw, linear2_b_tw, c_o, c_i);
    kutacc_af2_trans_act_inputs_t_wrapper *trans_inputs_ptr = new kutacc_af2_trans_act_inputs_t_wrapper(input_act_tw, intermediate_act_tw, batch, n_res);
    kutacc_af2_transition(trans_inputs_ptr, trans_weights_ptr, out_tw.get_tensor());
    delete trans_weights_ptr;
    delete trans_inputs_ptr;
    return out;
}
}

// bind.h
#include <torch/extension.h>
#include "test_transition.h"

namespace alphafold {
inline void bind(pybind11::module &m)
{
    auto submodule = m.def_submodule("alphafold");
    submodule.def("test_transition", &test_transition, py::arg("batch"), py::arg("n_res"), 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_transition(batch, n_res, c_o, c_i):
    out = kernel.alphafold.test_transition(batch, n_res, c_o, c_i)
    return out

Parent topic: AlphaFold2