Commit fc1f5871 authored by Jaikrishnan Menon's avatar Jaikrishnan Menon

CPU Direct Execution: Move MatMulBias builder to a separate TU

parent 605eb460
/*******************************************************************************
* Copyright 2018 Intel Corporation
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*******************************************************************************/
#include "ngraph/runtime/cpu/op/matmul_bias.hpp"
#include "ngraph/runtime/cpu/cpu_builder.hpp"
#include "ngraph/runtime/cpu/cpu_kernels.hpp"
using namespace std;
using namespace ngraph;
namespace ngraph
{
namespace runtime
{
namespace cpu
{
template <>
void Builder::BUILDER_DECL(ngraph::op::MatmulBias)
{
auto& functors = external_function->get_functors();
auto& tensor_data = external_function->get_tensor_data();
auto& arg0_tensor = tensor_data[args[0].get_name()];
auto& arg1_tensor = tensor_data[args[1].get_name()];
auto& out0_tensor = tensor_data[out[0].get_name()];
const ngraph::op::MatmulBias* mm = static_cast<const ngraph::op::MatmulBias*>(node);
const auto& arg0_shape = mm->get_arg0_shape();
const auto& arg1_shape = mm->get_arg1_shape();
const auto& arg2_shape = node->get_shape();
auto m = arg0_shape[0];
auto n = arg1_shape[1];
auto k = arg0_shape[1];
bool transpose_A = false, transpose_B = false;
auto lda = arg0_shape[1];
auto ldb = arg1_shape[1];
if (mm->get_is_arg0_transposed())
{
transpose_A = true;
m = arg0_shape[1];
k = arg0_shape[0];
}
if (mm->get_is_arg1_transposed())
{
transpose_B = true;
n = arg1_shape[0];
}
const float beta = 0.0f;
auto mm_functor =
[&, transpose_A, transpose_B, m, n, k, lda, ldb, beta, arg2_shape](
CPURuntimeContext* ctx) {
cblas::cblas_sgemm(
cblas::Layout::RowMajor,
transpose_A ? cblas::Transpose::Transpose : cblas::Transpose::None,
transpose_B ? cblas::Transpose::Transpose : cblas::Transpose::None,
m,
n,
k,
1.0f,
static_cast<float*>(arg0_tensor),
max(1UL, lda),
static_cast<float*>(arg1_tensor),
max(1UL, ldb),
beta,
static_cast<float*>(out0_tensor),
max(1UL, arg2_shape[1]));
};
function<void(CPURuntimeContext*)> bias_functor = [](CPURuntimeContext* ctx) {};
if (args.size() > 2)
{
auto& arg2_tensor = tensor_data[args[2].get_name()];
auto axes = mm->get_broadcast_axes();
if (axes.size() == 1)
{
if (*(axes.begin()) == 0)
{
vector<float> ones_row(arg2_shape[0], 1.0f);
bias_functor = [&, ones_row, arg2_shape](CPURuntimeContext* ctx) {
cblas::cblas_sgemm(cblas::Layout::RowMajor,
cblas::Transpose::None,
cblas::Transpose::None,
arg2_shape[0],
arg2_shape[1],
1,
1.0f,
ones_row.data(),
1UL,
static_cast<float*>(arg2_tensor),
max(1UL, arg2_shape[1]),
1.0f,
static_cast<float*>(out0_tensor),
max(1UL, arg2_shape[1]));
};
}
else
{
vector<float> ones_col(arg2_shape[1], 1.0f);
bias_functor = [&, ones_col, arg2_shape](CPURuntimeContext* ctx) {
cblas::cblas_sgemm(cblas::Layout::RowMajor,
cblas::Transpose::None,
cblas::Transpose::None,
arg2_shape[0],
arg2_shape[1],
1,
1.0f,
static_cast<float*>(arg2_tensor),
1UL,
ones_col.data(),
max(1UL, arg2_shape[1]),
1.0f,
static_cast<float*>(out0_tensor),
max(1UL, arg2_shape[1]));
};
}
}
else
{
if (axes.size() != 2)
{
throw ngraph_error("unexpected broadcast rank");
}
vector<float> ones_scalar(arg2_shape[0], 1.0f);
bias_functor = [&, ones_scalar, arg2_shape](CPURuntimeContext* ctx) {
vector<float> bias(arg2_shape[1], *static_cast<float*>(arg2_tensor));
cblas::cblas_sgemm(cblas::Layout::RowMajor,
cblas::Transpose::None,
cblas::Transpose::None,
arg2_shape[0],
arg2_shape[1],
1,
1.0f,
ones_scalar.data(),
1UL,
bias.data(),
max(1UL, arg2_shape[1]),
1.0f,
static_cast<float*>(out0_tensor),
max(1UL, arg2_shape[1]));
};
}
}
auto functor = [&, mm_functor, bias_functor](CPURuntimeContext* ctx) {
mm_functor(ctx);
bias_functor(ctx);
};
functors.emplace_back(functor);
}
REGISTER_OP_BUILDER(MatmulBias);
}
}
}
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment