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 * \brief External function interface to rocBLAS libraries
22 * \file tags.h
23 */
24 #ifndef TOPI_CONTRIB_ROCBLAS_H_
25 #define TOPI_CONTRIB_ROCBLAS_H_
26
27 #include "tvm/operation.h"
28 #include "topi/detail/extern.h"
29
30 namespace topi {
31 namespace contrib {
32 using namespace tvm;
33 /*!
34 * \brief Create an op that multiplies lhs and rhs with rocBLAS
35 *
36 * \param lhs The left matrix operand
37 * \param rhs The right matrix operand
38 * \param transa Whether to transpose lhs
39 * \param transb Whether to transpose rhs
40 *
41 * \return The output tensor
42 */
rocblas_matmul(const Tensor & lhs,const Tensor & rhs,bool transa,bool transb)43 inline Tensor rocblas_matmul(const Tensor& lhs,
44 const Tensor& rhs,
45 bool transa,
46 bool transb) {
47 auto n = transa ? lhs->shape[1] : lhs->shape[0];
48 auto m = transb ? rhs->shape[0] : rhs->shape[1];
49
50 return make_extern(
51 { { n, m } }, { lhs->dtype }, { lhs, rhs },
52 [&](Array<Buffer> ins, Array<Buffer> outs) {
53 return call_packed({
54 Expr("tvm.contrib.rocblas.matmul"),
55 pack_buffer(ins[0]),
56 pack_buffer(ins[1]),
57 pack_buffer(outs[0]),
58 transa,
59 transb });
60 }, "C", "", {})[0];
61 }
62
63 } // namespace contrib
64 } // namespace topi
65
66 #endif // TOPI_CONTRIB_ROCBLAS_H_
67