exec_context.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */
#ifndef TVM_TIRX_EXEC_CONTEXT_H_
#define TVM_TIRX_EXEC_CONTEXT_H_
 
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/layout.h>
#include <tvm/tirx/var.h>
 
#include <string>
#include <unordered_map>
#include <vector>
 
namespace tvm {
namespace tirx {
 
constexpr int kWgSize = 4;
 
struct AxisRange {
  PrimExpr extent;
  PrimExpr offset;
  PrimExpr stride;
 
  bool Intersect(int64_t lo, int64_t hi, AxisRange* out) const;
 
  bool Modulo(int64_t modulus, int64_t residue, AxisRange* out) const;
};
 
struct ActiveSet {
  TileLayout layout;
 
  int64_t size() const;
  bool GetAxis(const std::string& axis, AxisRange* out) const;
  bool HasAxis(const std::string& axis) const;
  ActiveSet WithAxis(const std::string& axis, const AxisRange& range) const;
  std::vector<std::string> AxisNames() const;
};
 
struct ExecSplit {
  std::unordered_map<std::string, AxisRange> inter;
  std::unordered_map<std::string, AxisRange> intra;
};
 
TVM_DLL ActiveSet InitialActiveSet(int64_t lane_ext, int64_t warp_ext, int64_t cta_ext);
TVM_DLL ActiveSet InitialActiveSet(int64_t lane_ext, int64_t warp_ext, int64_t cta_ext,
                                   const std::vector<std::pair<std::string, int64_t>>& cta_axes);
 
TVM_DLL bool FilterNarrow(const ActiveSet& A, ScopeBinding binding, int64_t lo, int64_t hi,
                          ActiveSet* out, std::string* err);
 
TVM_DLL bool ScopeSwitch(const ActiveSet& A, ScopeKind scope_kind, ExecSplit* out,
                         std::string* err);
 
struct ExecContext {
  ActiveSet A;
  ScopeKind scope_kind = ScopeKind::kThread;
  ExecSplit split;  // (inter, intra) of current A under current scope_kind
 
  static ExecContext AtKernelEntry(int64_t lane_ext, int64_t warp_ext, int64_t cta_ext);
  static ExecContext AtKernelEntry(int64_t lane_ext, int64_t warp_ext, int64_t cta_ext,
                                   const std::vector<std::pair<std::string, int64_t>>& cta_axes);
 
  bool WithFilter(ScopeBinding binding, int64_t lo, int64_t hi, ExecContext* out,
                  std::string* err) const;
 
  bool WithSelector(ScopeBinding binding, PrimExpr selector, ExecContext* out,
                    std::string* err) const;
 
  bool WithCtaAxisFilter(const std::string& axis, int64_t lo, int64_t hi, ExecContext* out,
                         std::string* err) const;
 
  bool WithCtaAxisModulo(const std::string& axis, int64_t modulus, int64_t residue,
                         ExecContext* out, std::string* err) const;
};
 
TVM_DLL ffi::Map<ffi::String, ffi::Array<PrimExpr>> EncodeSplitSide(
    const std::unordered_map<std::string, AxisRange>& side);
 
}  // namespace tirx
}  // namespace tvm
 
#endif  // TVM_TIRX_EXEC_CONTEXT_H_