"test/vscode:/vscode.git/clone" did not exist on "03b0364f7661035f0e94a11ea4925b1b47d55074"
gemm.h 1.02 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

/*!
 * \file tl/op/gemm.h
 * \brief Define gemm operator.
 *
 */

#ifndef TVM_TL_OP_GEMM_H_
#define TVM_TL_OP_GEMM_H_

#include "op.h"

namespace tvm {
namespace tl {

using namespace tir;

class Gemm : public Operator {
21
public:
22
  Gemm(Array<PrimExpr> args, BufferMap vmap);
23
24
25
  Stmt Lower(const LowerArgs &T, arith::Analyzer *analyzer) const final;
  LayoutMap InferLayout(const LayoutInferArgs &T, InferLevel level) final;
  static const Op &Get();
26
27
28
29
30
31
  enum class GemmWarpPolicy {
    kSquare = 0,
    kFullRow = 1,
    kFullCol = 2,
  } policy;

32
private:
33
34
35
36
37
38
39
40
  std::pair<int, int> ComputeWarpPartition(int num_warps, Target target) const;

  Array<PrimExpr> call_args;
  tir::Buffer A, B, C;
  bool trans_A, trans_B;
  int M, N, K;
  // k_pack please ref to bitblas/tl/mfma_macro_generator.py::k_pack
  // only will be enabled under cdna mfma instructions
41
  int kPack = 1;
42
43
44
  bool completed_ = false;
};

45
46
} // namespace tl
} // namespace tvm
47

48
#endif //  TVM_TL_OP_GEMM_H_