StructuredOpsUtils.h

Go to the documentation of this file.

//===- StructuredOpsUtils.h - Utilities used by structured ops --*- C++ -*-===//
//
// Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This header file define utilities that operate on standard types and are
// useful across multiple dialects that use structured ops abstractions. These
// abstractions consist of define custom operations that encode and transport
// information about their semantics (e.g. type of iterators like parallel,
// reduction, etc..) as attributes.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_DIALECT_UTILS_STRUCTUREDOPSUTILS_H
#define MLIR_DIALECT_UTILS_STRUCTUREDOPSUTILS_H

#include "mlir/IR/Attributes.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/StringRef.h"

namespace mlir {
static constexpr StringLiteral getIndexingMapsAttrName() {
  return StringLiteral("indexing_maps");
}

static constexpr StringLiteral getIteratorTypesAttrName() {
  return StringLiteral("iterator_types");
}

static constexpr StringLiteral getArgsInAttrName() {
  return StringLiteral("args_in");
}

static constexpr StringLiteral getArgsOutAttrName() {
  return StringLiteral("args_out");
}

static constexpr StringLiteral getDocAttrName() { return StringLiteral("doc"); }

static constexpr StringLiteral getFunAttrName() { return StringLiteral("fun"); }

static constexpr StringLiteral getLibraryCallAttrName() {
  return StringLiteral("library_call");
}

inline static constexpr StringLiteral getParallelIteratorTypeName() {
  return StringLiteral("parallel");
}

inline static constexpr StringLiteral getReductionIteratorTypeName() {
  return StringLiteral("reduction");
}

inline static constexpr StringLiteral getWindowIteratorTypeName() {
  return StringLiteral("window");
}

inline static ArrayRef<StringRef> getAllIteratorTypeNames() {
  static StringRef names[3] = {getParallelIteratorTypeName(),
                               getReductionIteratorTypeName(),
                               getWindowIteratorTypeName()};
  return llvm::makeArrayRef(names);
}

inline unsigned getNumIterators(StringRef name, ArrayAttr iteratorTypes) {
  auto names = getAllIteratorTypeNames();
  (void)names;
  assert(llvm::is_contained(names, name));
  return llvm::count_if(iteratorTypes, [name](Attribute a) {
    return a.cast<StringAttr>().getValue() == name;
  });
}

inline unsigned getNumIterators(ArrayAttr iteratorTypes) {
  unsigned res = 0;
  for (auto n : getAllIteratorTypeNames())
    res += getNumIterators(n, iteratorTypes);
  return res;
}

} // end namespace mlir

#endif // MLIR_UTILS_STRUCTUREDOPSUTILS_H