1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19
20/*!
21 * \file rocm/dense.h
22 * \brief rocm schedule for dense operation
23 */
24#ifndef TVM_TOPI_ROCM_DENSE_H_
25#define TVM_TOPI_ROCM_DENSE_H_
26
27#include <tvm/target/generic_func.h>
28#include <tvm/te/operation.h>
29#include <tvm/topi/contrib/rocblas.h>
30#include <tvm/topi/cuda/dense.h>
31#include <tvm/topi/detail/array_utils.h>
32#include <tvm/topi/generic/extern.h>
33#include <tvm/topi/nn/dense.h>
34#include <tvm/topi/tags.h>
35
36namespace tvm {
37namespace topi {
38
39using namespace tvm::te;
40
41namespace rocm {
42/*!
43 * \brief Implementation of dense for rocm backend
44 *
45 * \param target The target device
46 * \param data Tensor with shape [batch, in_dim]
47 * \param weight Tensor with shape [out_dim, in_dim]
48 * \param bias Tensor with shape [out_dim]. Optional; to omit bias, pass Tensor()
49 * \param out_dtype Output data type. Used for mixed precision.
50 *
51 * \return Tensor with shape [batch, out_dim]
52 */
53inline tvm::te::Tensor dense_rocm(const Target& target, const tvm::te::Tensor& data,
54 const tvm::te::Tensor& weight, const tvm::te::Tensor& bias,
55 const DataType& out_dtype) {
56 ICHECK_EQ(data->shape.size(), 2) << "dense requires 2-D data";
57 ICHECK_EQ(weight->shape.size(), 2) << "dense requires 2-D weight";
58 if (bias.defined()) {
59 ICHECK_EQ(bias->shape.size(), 1) << "dense requires 1-D bias";
60 }
61
62 auto batch = data->shape[0];
63 auto in_dim = data->shape[1];
64 auto out_dim = weight->shape[0];
65
66 if (target->GetLibs().count("rocblas")) {
67 ICHECK_EQ(data->dtype, out_dtype) << "Mixed precision not supported.";
68 auto mm = topi::contrib::rocblas_matmul(data, weight, false, true);
69 if (bias.defined()) {
70 mm = tvm::te::compute(
71 {batch, out_dim}, [&](Var i, Var j) { return mm(i, j) + bias(j); }, "tensor", kBroadcast);
72 }
73
74 return mm;
75 } else {
76 return topi::nn::dense(data, weight, bias, out_dtype);
77 }
78}
79
80/*!
81 * \brief Create a rocm schedule for dense
82 *
83 * \param target The target to generate a schedule for.
84 * \param outs The output tensors.
85 *
86 * \return A schedule for the given ops.
87 */
88inline Schedule schedule_dense(const Target& target, const Array<Tensor>& outs) {
89 if (target->kind->name == "rocm" && target->GetLibs().count("rocblas")) {
90 return topi::generic::schedule_extern(target, outs);
91 }
92
93 return topi::cuda::schedule_dense(target, outs);
94}
95
96} // namespace rocm
97} // namespace topi
98} // namespace tvm
99#endif // TVM_TOPI_ROCM_DENSE_H_
100