maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

Implement simple type polymorphism for linalg named ops. 

* It was decided that this was the end of the line for the existing custom tc parser/generator, and this is the first step to replacing it with a declarative format that maps well to mathy source languages.
* One such source language is implemented here: https://github.com/stellaraccident/mlir-linalgpy/blob/main/samples/mm.py
  * In fact, this is the exact source of the declarative `polymorphic_matmul` in this change.
  * I am working separately to clean this python implementation up and add it to MLIR (probably as `mlir.tools.linalg_opgen` or equiv). The scope of the python side is greater than just generating named ops: the ops are callable and directly emit `linalg.generic` ops fully dynamically, and this is intended to be a feature for frontends like npcomp to define custom linear algebra ops at runtime.
* There is more work required to handle full type polymorphism, especially with respect to integer formulations, since they require more specificity wrt types.
* Followups to this change will bring the new generator to feature parity with the current one and delete the current. Roughly, this involves adding support for interface declarations and attribute symbol bindings.

Differential Revision: https://reviews.llvm.org/D97135
This commit is contained in:
Stella Laurenzo 2021-02-20 22:46:27 -08:00
parent b568d3d6c9
commit 6c9541d4dd
9 changed files with 1131 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
mlir/docs/Dialects/Linalg.md
View File

@ -662,6 +662,18 @@ void batchmatmul::regionBuilder(ArrayRef<BlockArgument> args) {
}
```
### YAML Based Named Structured Ops
Linalg provides a declarative generation tool (`mlir-linalg-ods-yaml-gen`) to
automatically produce named ops from a YAML-based op description format
intended to capture the structure of the named ops and be generated from a
higher level "mathy" DSL syntax. This facility is currently in flight and is
intended to subsume the above when ready. See the C++ class to YAML mapping
traits in `mlir-mlinalg-ods-yaml-gen.cpp` as the source of truth for the schema.
Most of the above documentation roughly applies to this path and will be ported
as migration continues.
## Open Issues and Design Alternatives<a name="open_issues"></a>
Multiple open issues and design alternatives are in flight and it is time to lay

47
mlir/include/mlir/Dialect/Linalg/IR/CMakeLists.txt
View File

@ -1,8 +1,8 @@
# Declare a function to generate ODS with mlir-linalg-ods-gen
function(add_linalg_ods_gen tc_filename output_file)
function(add_linalg_ods_tc_gen tc_filename output_file)
set(TC_SOURCE ${CMAKE_CURRENT_SOURCE_DIR}/${tc_filename})
set(GEN_ODS_FILE ${CMAKE_CURRENT_BINARY_DIR}/${output_file}.td)
set(GEN_CPP_FILE ${CMAKE_CURRENT_BINARY_DIR}/${output_file}.cpp.inc)
set(GEN_ODS_FILE ${CMAKE_CURRENT_BINARY_DIR}/${output_file}.tcgen.td)
set(GEN_CPP_FILE ${CMAKE_CURRENT_BINARY_DIR}/${output_file}.tcgen.cpp.inc)
set_source_files_properties(
${GEN_ODS_FILE}
PROPERTIES GENERATED TRUE)
@ -20,17 +20,52 @@ function(add_linalg_ods_gen tc_filename output_file)
${MLIR_LINALG_ODS_GEN_TARGET}
VERBATIM)
add_custom_target(
MLIR${output_file}IncGen
MLIR${output_file}TcIncGen
DEPENDS
${MLIR_LINALG_ODS_GEN_EXE}
${MLIR_LINALG_ODS_GEN_TARGET}
${GEN_ODS_FILE} ${GEN_CPP_FILE})
endfunction()
add_linalg_ods_gen(LinalgNamedStructuredOpsSpec.tc LinalgNamedStructuredOps)
# Declare a function to generate ODS with mlir-linalg-ods-yaml-gen
function(add_linalg_ods_yaml_gen yaml_ast_file output_file)
set(YAML_AST_SOURCE ${CMAKE_CURRENT_SOURCE_DIR}/${yaml_ast_file})
set(GEN_ODS_FILE ${CMAKE_CURRENT_BINARY_DIR}/${output_file}.yamlgen.td)
set(GEN_CPP_FILE ${CMAKE_CURRENT_BINARY_DIR}/${output_file}.yamlgen.cpp.inc)
set_source_files_properties(
${GEN_ODS_FILE}
PROPERTIES GENERATED TRUE)
set_source_files_properties(
${GEN_CPP_FILE}
PROPERTIES GENERATED TRUE)
add_custom_command(
OUTPUT ${GEN_ODS_FILE} ${GEN_CPP_FILE}
COMMAND ${MLIR_LINALG_ODS_YAML_GEN_EXE} ${YAML_AST_SOURCE} -o-ods-decl=${GEN_ODS_FILE} -o-impl=${GEN_CPP_FILE}
MAIN_DEPENDENCY
${YAML_AST_SOURCE}
DEPENDS
${MLIR_LINALG_ODS_YAML_GEN_EXE}
${MLIR_LINALG_ODS_YAML_GEN_TARGET})
add_custom_target(
MLIR${output_file}YamlIncGen
DEPENDS
${MLIR_LINALG_ODS_YAML_GEN_EXE}
${MLIR_LINALG_ODS_YAML_GEN_TARGET}
${GEN_ODS_FILE} ${GEN_CPP_FILE})
endfunction()
# TODO: Delete tc generation and replace with the YAML variant once all ops are
# ported.
add_linalg_ods_tc_gen(LinalgNamedStructuredOpsSpec.tc LinalgNamedStructuredOps)
add_linalg_ods_yaml_gen(LinalgNamedStructuredOps.yaml LinalgNamedStructuredOps)
# Provide a short name for all external dependency that needs to
# include Linalg in ODS
add_custom_target(LinalgOdsGen DEPENDS MLIRLinalgNamedStructuredOpsIncGen)
add_custom_target(LinalgOdsGen
DEPENDS
MLIRLinalgNamedStructuredOpsTcIncGen
MLIRLinalgNamedStructuredOpsYamlIncGen
)
add_dependencies(mlir-headers LinalgOdsGen)
add_mlir_dialect(LinalgOps linalg)

50
mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml Normal file
View File

@ -0,0 +1,50 @@
--- !LinalgOpConfig
metadata: !LinalgOpMetadata
name: polymorphic_matmul
cpp_op_name: PolymorphicMatmulOp
doc: |-
Type polymorphic matrix multiplication.
This op is presently here to test a new path for generation and will replace
the existing 'matmul' op when ready. Do not use.
structured_op: !LinalgStructuredOpConfig
args:
- !<LinalgTensorDef>
name: A
usage: input
shape: affine_map<()[s0, s1, s2] -> (s0, s2)>
- !<LinalgTensorDef>
name: B
usage: input
shape: affine_map<()[s0, s1, s2] -> (s2, s1)>
- !<LinalgTensorDef>
name: C
usage: output
shape: affine_map<()[s0, s1, s2] -> (s0, s1)>
indexing_maps: !LinalgIndexingMapsConfig
static_indexing_maps:
- affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d2)>
- affine_map<(d0, d1, d2)[s0, s1, s2] -> (d2, d1)>
- affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d1)>
iterator_types:
- parallel
- parallel
- reduction
assignments:
- !ScalarAssign
arg: C
value: !ScalarExpression
scalar_apply:
fn_name: add
operands:
- !ScalarExpression
scalar_arg: C
- !ScalarExpression
scalar_apply:
fn_name: mul
operands:
- !ScalarExpression
scalar_arg: A
- !ScalarExpression
scalar_arg: B

3
mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
View File

@ -806,6 +806,7 @@ def IndexedGenericOp : GenericOpBase<"indexed_generic"> {
//===----------------------------------------------------------------------===//
// This file is auto-generated from a TC def specification.
include "mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.td"
include "mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.tcgen.td"
include "mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yamlgen.td"
#endif // LINALG_STRUCTURED_OPS

1
mlir/lib/Dialect/Linalg/IR/CMakeLists.txt
View File

@ -14,6 +14,7 @@ add_mlir_dialect_library(MLIRLinalg
LINK_LIBS PUBLIC
MLIRAffine
MLIRIR
MLIRParser
MLIRSideEffectInterfaces
MLIRViewLikeInterface
MLIRStandard

82
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
View File

@ -20,6 +20,7 @@
#include "mlir/IR/Matchers.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Parser.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SetVector.h"
@ -121,6 +122,81 @@ static LogicalResult foldMemRefCast(Operation *op) {
return success(folded);
}
//===----------------------------------------------------------------------===//
// Region builder helper.
// TODO: Move this to a utility library.
// The public methods on this class are referenced directly from generated code
// and bind by name to math functions in the DSL as:
// `applyfn__{fnName}`
// Examples:
// `applyfn__add`
// `applyfn__mul`
// The naming convention is intentional in order to match snake-cased DSL names.
// See mlir-linalg-ods-yaml-gen.cpp for the code that mates to this class.
//
// Implementations of the math functions must be polymorphic over numeric types,
// internally performing necessary casts. If the function application makes no
// sense, then the only recourse is to assert and return nullptr. This can be
// extended later if it becomes possible to fail construction of the region. The
// invariant should be enforced at a higher level.
//
// TODO: These helpers are currently type polymorphic over the class of integer
// and floating point types, but they will not internally cast within bit
// widths of a class (mixed precision such as i8->i32) or across classes
// (i.e. mixed float and integer). Many such combinations are ambiguous or need
// to be handled with care and work is being considered to extend the op
// language to make such cases explicit. In the mean-time, violating this will
// fail verification, which is deemed acceptable.
//===----------------------------------------------------------------------===//
namespace {
class RegionBuilderHelper {
public:
RegionBuilderHelper(Block &block) : block(block) {}
Value applyfn__add(Value lhs, Value rhs) {
OpBuilder builder = getBuilder(lhs);
if (isFloatingPoint(lhs))
return builder.create<AddFOp>(lhs.getLoc(), lhs, rhs);
else if (isInteger(lhs))
return builder.create<AddIOp>(lhs.getLoc(), lhs, rhs);
llvm_unreachable("unsupported non numeric type");
}
Value applyfn__mul(Value lhs, Value rhs) {
OpBuilder builder = getBuilder(lhs);
if (isFloatingPoint(lhs))
return builder.create<MulFOp>(lhs.getLoc(), lhs, rhs);
else if (isInteger(lhs))
return builder.create<MulIOp>(lhs.getLoc(), lhs, rhs);
llvm_unreachable("unsupported non numeric type");
}
void yieldOutputs(ValueRange values) {
assert(!values.empty() && "linalg ops must yield outputs");
if (values.empty())
return;
Value first = values.front();
OpBuilder builder = getBuilder(first);
builder.create<YieldOp>(first.getLoc(), values);
}
private:
Block &block;
bool isFloatingPoint(Value value) { return value.getType().isa<FloatType>(); }
bool isInteger(Value value) { return value.getType().isa<IntegerType>(); }
OpBuilder getBuilder(Value value) {
OpBuilder builder(value.getContext());
builder.setInsertionPointToEnd(&block);
return builder;
}
};
} // namespace
//===----------------------------------------------------------------------===//
// CopyOp
//===----------------------------------------------------------------------===//
@ -1868,7 +1944,8 @@ struct EraseDeadLinalgOp;
struct FoldTensorCastOp;
} // namespace
#include "mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.cpp.inc"
#include "mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.tcgen.cpp.inc"
#include "mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yamlgen.cpp.inc"
#define GET_OP_CLASSES
#include "mlir/Dialect/Linalg/IR/LinalgOps.cpp.inc"
@ -2032,7 +2109,8 @@ fillStructuredOpRegion(OpBuilder &opBuilder, Region &region,
unsigned actual = body->getNumArguments();
unsigned expected = NamedStructuredOpType::getNumRegionArgs();
if (expected != actual) {
if (errorHandler) errorHandler(expected, actual);
if (errorHandler)
errorHandler(expected, actual);
return;
}

27
mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir Normal file
View File

@ -0,0 +1,27 @@
// RUN: mlir-opt %s -split-input-file -linalg-generalize-named-ops | FileCheck %s
func @generalize_matmul_tensor_f32(%A : tensor<16x8xf32>, %B: tensor<8x32xf32>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> {
%0 = linalg.polymorphic_matmul ins(%A, %B: tensor<16x8xf32>, tensor<8x32xf32>)
outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
return %0: tensor<16x32xf32>
}
// CHECK: ^{{.*}}(%[[A_ARG:.+]]: f32, %[[B_ARG:.+]]: f32, %[[C_ARG:.+]]: f32)
// CHECK-NEXT: %[[MUL:.+]] = mulf %[[A_ARG]], %[[B_ARG]] : f32
// CHECK-NEXT: %[[ADD:.+]] = addf %[[C_ARG]], %[[MUL]] : f32
// CHECK-NEXT: linalg.yield %[[ADD]] : f32
// CHECK-NEXT: -> tensor<16x32xf32>
// -----
func @generalize_matmul_tensor_i32(%A : tensor<16x8xi32>, %B: tensor<8x32xi32>, %C: tensor<16x32xi32>) -> tensor<16x32xi32> {
%0 = linalg.polymorphic_matmul ins(%A, %B: tensor<16x8xi32>, tensor<8x32xi32>)
outs(%C: tensor<16x32xi32>) -> tensor<16x32xi32>
return %0: tensor<16x32xi32>
}
// CHECK: ^{{.*}}(%[[A_ARG:.+]]: i32, %[[B_ARG:.+]]: i32, %[[C_ARG:.+]]: i32)
// CHECK-NEXT: %[[MUL:.+]] = muli %[[A_ARG]], %[[B_ARG]] : i32
// CHECK-NEXT: %[[ADD:.+]] = addi %[[C_ARG]], %[[MUL]] : i32
// CHECK-NEXT: linalg.yield %[[ADD]] : i32
// CHECK-NEXT: -> tensor<16x32xi32>

39
mlir/tools/mlir-linalg-ods-gen/CMakeLists.txt
View File

@ -2,6 +2,13 @@ set(LLVM_LINK_COMPONENTS
Core
Support
)
set(LLVM_OPTIONAL_SOURCES
mlir-linalg-ods-gen.cpp
mlir-linalg-ods-yaml-gen.cpp
)
# Original mlir-linalg-ods-gen (to be replaced).
add_llvm_tool(mlir-linalg-ods-gen
mlir-linalg-ods-gen.cpp
)
@ -30,3 +37,35 @@ if(LLVM_USE_HOST_TOOLS)
endif()
endif()
endif()
# New mlir-linalg-ods-yaml-gen.
add_llvm_tool(mlir-linalg-ods-yaml-gen
mlir-linalg-ods-yaml-gen.cpp
)
llvm_update_compile_flags(mlir-linalg-ods-yaml-gen)
target_link_libraries(mlir-linalg-ods-yaml-gen PRIVATE
MLIRIR
MLIRSupport
MLIRParser
)
set(MLIR_LINALG_ODS_YAML_GEN mlir-linalg-ods-yaml-gen CACHE
STRING "Native mlir-linalg-ods-yaml-gen executable. Saves building one when cross-compiling.")
set(MLIR_LINALG_ODS_YAML_GEN_EXE ${MLIR_LINALG_ODS_YAML_GEN} PARENT_SCOPE)
set(MLIR_LINALG_ODS_YAML_GEN_TARGET mlir-linalg-ods-yaml-gen PARENT_SCOPE)
if(LLVM_USE_HOST_TOOLS)
if ("${MLIR_LINALG_ODS_YAML_GEN_EXE}" STREQUAL mlir-linalg-ods-yaml-gen)
build_native_tool(mlir-linalg-ods-yaml-gen MLIR_LINALG_ODS_YAML_GEN_EXE DEPENDS mlir-linalg-ods-yaml-gen)
set(MLIR_LINALG_ODS_YAML_GEN_EXE ${MLIR_LINALG_ODS_YAML_GEN_EXE} PARENT_SCOPE)
add_custom_target(mlir-linalg-ods-yaml-gen-host DEPENDS ${MLIR_LINALG_ODS_YAML_GEN_EXE})
set(MLIR_LINALG_ODS_YAML_GEN_TARGET mlir-linalg-ods-yaml-gen-host DEPENDS PARENT_SCOPE)
if(NOT LLVM_BUILD_UTILS)
set_target_properties(mlir-linalg-ods-yaml-gen PROPERTIES EXCLUDE_FROM_ALL ON)
endif()
endif()
endif()

878
mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp Normal file
View File

@ -0,0 +1,878 @@
//===- mlir-linalg-ods-yaml-gen.cpp - Linalg ODS generation from yaml ----===//
//
// Part of the LLVM 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 file implements an ODS (and C++) generator from a YAML form
// derived from the mathematical expression of linalg named ops. Typically a
// math oriented DSL will be used to export the essential representation to
// this form, and maintaining the SOT at the math level (versus recreating it
// in MLIR) is deemed to have systemic value.
//
//===----------------------------------------------------------------------===//
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/Parser.h"
#include "mlir/Support/FileUtilities.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/Optional.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/FormatVariadic.h"
#include "llvm/Support/ToolOutputFile.h"
#include "llvm/Support/YAMLTraits.h"
using namespace mlir;
using llvm::yaml::Input;
using llvm::yaml::IO;
using llvm::yaml::MappingTraits;
using llvm::yaml::ScalarEnumerationTraits;
using llvm::yaml::ScalarTraits;
#define DEBUG_TYPE "linalg-ods-gen"
//===----------------------------------------------------------------------===//
// Mapping structs (correspond to data types in the YAML description).
// TODO: Since this is a schema/part of the contract, it should be moved to
// a real header.
//===----------------------------------------------------------------------===//
namespace {
struct LinalgYAMLContext {
MLIRContext *mlirContext;
};
struct LinalgOpMetadata {
std::string name;
std::string cppOpName;
Optional<std::string> doc;
};
struct SerializedAffineMap {
AffineMapAttr affineMapAttr;
AffineMap affineMap() { return affineMapAttr.getValue(); }
};
enum class LinalgTensorUsageDef {
input,
output,
temporary,
};
struct LinalgTensorDef {
std::string name;
LinalgTensorUsageDef usage;
SerializedAffineMap shape;
};
enum class LinalgIteratorTypeDef {
parallel,
reduction,
};
struct LinalgIndexingMapsConfig {
Optional<SmallVector<SerializedAffineMap>> staticIndexingMaps;
};
struct ScalarExpression;
struct ScalarApply {
std::string fnName;
// NOTE: Must be pure heap allocated container (not SmallVector)
// due to recursive data type.
std::vector<ScalarExpression> operands;
};
struct ScalarExpression {
Optional<std::string> scalarArg;
Optional<ScalarApply> scalarApply;
};
struct ScalarAssign {
std::string arg;
ScalarExpression value;
};
struct LinalgStructuredOpConfig {
SmallVector<LinalgTensorDef> args;
LinalgIndexingMapsConfig indexingMaps;
SmallVector<LinalgIteratorTypeDef> iteratorTypes;
SmallVector<ScalarAssign> assignments;
};
struct LinalgOpConfig {
Optional<LinalgOpMetadata> metadata;
Optional<LinalgStructuredOpConfig> structuredOp;
};
} // namespace
//===----------------------------------------------------------------------===//
// Mapping traits.
//===----------------------------------------------------------------------===//
LLVM_YAML_IS_SEQUENCE_VECTOR(LinalgTensorDef);
LLVM_YAML_IS_SEQUENCE_VECTOR(SerializedAffineMap);
LLVM_YAML_IS_SEQUENCE_VECTOR(LinalgIteratorTypeDef);
LLVM_YAML_IS_SEQUENCE_VECTOR(ScalarAssign);
LLVM_YAML_IS_SEQUENCE_VECTOR(ScalarExpression);
LLVM_YAML_IS_DOCUMENT_LIST_VECTOR(LinalgOpConfig);
namespace llvm {
namespace yaml {
/// Top-level type containing op metadata and one of a concrete op type.
/// Currently, the only defined op type is `structured_op` (maps to
/// `LinalgStructuredOpConfig`).
template <>
struct MappingTraits<LinalgOpConfig> {
static void mapping(IO &io, LinalgOpConfig &info) {
io.mapOptional("metadata", info.metadata);
io.mapOptional("structured_op", info.structuredOp);
}
};
/// A structured op models (at most) a single contraction by modeling
/// - A list of named arguments (`LinalgTensorDef`), which can be inputs,
/// outputs, or temporaries.
/// - List of indexing maps (see `LinalgIndexingMaps`).
/// - Iterator types (see `LinalgIteratorTypeDef`).
/// - List of scalar level assignment (see `ScalarAssign`).
template <>
struct MappingTraits<LinalgStructuredOpConfig> {
static void mapping(IO &io, LinalgStructuredOpConfig &info) {
io.mapRequired("args", info.args);
io.mapRequired("indexing_maps", info.indexingMaps);
io.mapRequired("iterator_types", info.iteratorTypes);
io.mapRequired("assignments", info.assignments);
}
};
/// Maps a named tensor-argument to an operation, consisting of:
/// - `name`: Must be unique within the operation.
/// - `usage`: How the argument is used (input, output, etc).
/// - `shape`: An AffineMap from all op symbols to the specific shape
/// of this argument. Each shape must be normalized over the same list of
/// symbols and have no dimension inputs.
template <>
struct MappingTraits<LinalgTensorDef> {
static void mapping(IO &io, LinalgTensorDef &info) {
io.mapRequired("name", info.name);
io.mapRequired("usage", info.usage);
io.mapRequired("shape", info.shape);
}
};
/// Usage enum for a named argument.
template <>
struct ScalarEnumerationTraits<LinalgTensorUsageDef> {
static void enumeration(IO &io, LinalgTensorUsageDef &value) {
io.enumCase(value, "input", LinalgTensorUsageDef::input);
io.enumCase(value, "output", LinalgTensorUsageDef::output);
io.enumCase(value, "temporary", LinalgTensorUsageDef::temporary);
}
};
/// Iterator type enum.
template <>
struct ScalarEnumerationTraits<LinalgIteratorTypeDef> {
static void enumeration(IO &io, LinalgIteratorTypeDef &value) {
io.enumCase(value, "parallel", LinalgIteratorTypeDef::parallel);
io.enumCase(value, "reduction", LinalgIteratorTypeDef::reduction);
}
};
/// Metadata about the op (name, C++ name, and documentation).
template <>
struct MappingTraits<LinalgOpMetadata> {
static void mapping(IO &io, LinalgOpMetadata &info) {
io.mapRequired("name", info.name);
io.mapRequired("cpp_op_name", info.cppOpName);
io.mapOptional("doc", info.doc);
}
};
/// How the ops indexing maps are produced. Must be one of:
/// - static_indexing_maps: A static list of AffineMaps, possibly with
/// some symbols that bind to attributes of the op. Each indexing map must
/// be normalized over the same list of dimensions, and its symbols must
/// match the symbols for argument shapes.
template <>
struct MappingTraits<LinalgIndexingMapsConfig> {
static void mapping(IO &io, LinalgIndexingMapsConfig &info) {
io.mapOptional("static_indexing_maps", info.staticIndexingMaps);
}
};
/// Models an assignment to a named output.
/// - The `arg` name must match a named output or temporary.
/// - The `value` is a scalar expression for computing the value to
/// assign (see `ScalarExpression`).
template <>
struct MappingTraits<ScalarAssign> {
static void mapping(IO &io, ScalarAssign &info) {
io.mapRequired("arg", info.arg);
io.mapRequired("value", info.value);
}
};
/// A scalar expression (RHS of an assignment). Must be one of:
/// - `scalar_arg`: Name of an argument to the op.
/// - `scalar_apply`: Result of evaluating a named function (see
/// `ScalarApply`).
template <>
struct MappingTraits<ScalarExpression> {
static void mapping(IO &io, ScalarExpression &info) {
io.mapOptional("scalar_arg", info.scalarArg);
io.mapOptional("scalar_apply", info.scalarApply);
}
};
/// A scalar expression that evaluates a named function.
/// Functions are generally "math" level and type polymorphic. Builtin
/// functions include:
/// - `add(lhs, rhs)`
/// - `mul(lhs, rhs)`
template <>
struct MappingTraits<ScalarApply> {
static void mapping(IO &io, ScalarApply &info) {
io.mapRequired("fn_name", info.fnName);
io.mapRequired("operands", info.operands);
}
};
/// Helper mapping which accesses an AffineMapAttr as a serialized string of
/// the same.
template <>
struct ScalarTraits<SerializedAffineMap> {
static void output(const SerializedAffineMap &value, void *rawYamlContext,
raw_ostream &out) {
assert(value.affineMapAttr);
value.affineMapAttr.print(out);
}
static StringRef input(StringRef scalar, void *rawYamlContext,
SerializedAffineMap &value) {
assert(rawYamlContext);
auto *yamlContext = static_cast<LinalgYAMLContext *>(rawYamlContext);
if (auto attr = mlir::parseAttribute(scalar, yamlContext->mlirContext)
.dyn_cast_or_null<AffineMapAttr>())
value.affineMapAttr = attr;
else if (!value.affineMapAttr || !value.affineMapAttr.isa<AffineMapAttr>())
return "could not parse as an affine map attribute";
return StringRef();
}
static QuotingType mustQuote(StringRef) { return QuotingType::None; }
};
} // namespace yaml
} // namespace llvm
namespace {
//===----------------------------------------------------------------------===//
// Generation utilities
//===----------------------------------------------------------------------===//
class GenerationContext {
public:
GenerationContext(MLIRContext *context, raw_ostream *odsOut,
raw_ostream *defnOut)
: context(context), loc(UnknownLoc::get(context)), odsOut(odsOut),
defnOut(defnOut) {}
MLIRContext *getContext() { return context; }
void setLoc(Location loc) { this->loc = loc; }
Location getLoc() { return loc; }
bool shouldGenerateOds() { return odsOut; }
bool shouldGenerateDefns() { return defnOut; }
raw_ostream &odss() {
assert(odsOut && "ODS stream not defined");
return *odsOut;
}
raw_ostream &defns() {
assert(defnOut && "Definition stream not defined");
return *defnOut;
}
private:
MLIRContext *context;
Location loc;
raw_ostream *odsOut;
raw_ostream *defnOut;
};
} // namespace
static std::string generateCppExpression(SerializedAffineMap self,
StringRef contextName) {
std::string printedStr;
llvm::raw_string_ostream printedSs(printedStr);
self.affineMapAttr.print(printedSs);
printedSs.flush();
static const char exprFormat[] =
R"FMT(mlir::parseAttribute("{0}", {1}).cast<AffineMapAttr>().getValue())FMT";
return llvm::formatv(exprFormat, printedStr, contextName);
}
template <typename Container>
static std::string interleaveToString(Container &container,
StringRef separator) {
std::string result;
llvm::raw_string_ostream ss(result);
llvm::interleave(container, ss, separator);
ss.flush();
return result;
}
static Optional<int>
findTensorDefArgIndex(StringRef name, SmallVectorImpl<LinalgTensorDef> &args) {
for (auto it : llvm::enumerate(args)) {
if (it.value().name == name)
return it.index();
}
return None;
}
static ScalarAssign *
findAssignment(StringRef name, SmallVectorImpl<ScalarAssign> &assignments) {
for (auto &assign : assignments) {
if (assign.arg == name)
return &assign;
}
return nullptr;
}
//===----------------------------------------------------------------------===//
// Templates
//===----------------------------------------------------------------------===//
// A single line banner format. Parameters:
// {0}: Single line comment
static const char bannerFormat[] = R"FMT(
//===----------------------------------------------------------------------===//
// {0}
//===----------------------------------------------------------------------===//
)FMT";
//===----------------------------------------------------------------------===//
// Named generic op generation.
// These ops map at most a single contraction that complies with the limitations
// of a linalg.generic.
//===----------------------------------------------------------------------===//
// Template for Linalg named ops' ODS definitions. Parameters:
// {0}: ODS/C++ op name
// {1}: assembly op mnemonic
// {2}: op interface list
// {3}: documentation (summary + description)
// {4}: op attribute list
// {5}: the number of arguments for the op region
// {6}: builder methods taking standalone attribute parameters
// {7}: additional methods for attributes used by indexing maps
static const char structuredOpOdsHeaderFormat[] = R"FMT(
//===----------------------------------------------------------------------===//
// Op definition for {0}
//===----------------------------------------------------------------------===//
def {0} : LinalgStructuredBase_Op<"{1}", [
AttrSizedOperandSegments,
DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
SingleBlockImplicitTerminator<"YieldOp">
/*extraInterfaces=*/{2}]> {
{3}
let arguments = (ins
Variadic<AnyShaped>:$inputs,
Variadic<AnyShaped>:$outputs{4}
);
let results = (outs Variadic<AnyRankedTensor>:$result_tensors);
let regions = (region AnyRegion:$region);
let skipDefaultBuilders = 1;
let builders = [
OpBuilderDAG<
(ins "ValueRange":$inputs, "ValueRange":$outputs),
[{{
$_state.addOperands(inputs);
$_state.addOperands(outputs);
$_state.addAttribute(
"operand_segment_sizes",
$_builder.getI32VectorAttr({{
static_cast<int32_t>(inputs.size()),
static_cast<int32_t>(outputs.size())}));
createAndFillStructuredOpRegion<{0}>(
$_builder,
$_state,
TypeRange(inputs),
TypeRange(outputs)/*, TODO: support captures*/);
}]>,
OpBuilderDAG<
(ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
"ValueRange":$outputs),
[{{
$_state.addOperands(inputs);
$_state.addOperands(outputs);
$_state.addTypes(resultTensorTypes);
$_state.addAttribute(
"operand_segment_sizes",
$_builder.getI32VectorAttr({{
static_cast<int32_t>(inputs.size()),
static_cast<int32_t>(outputs.size())}));
createAndFillStructuredOpRegion<{0}>(
$_builder,
$_state,
TypeRange(inputs),
TypeRange(outputs)/*, TODO: support captures*/);
}]>,
OpBuilderDAG<
(ins "TypeRange":$resultTensorTypes, "ValueRange":$operands,
CArg<"ArrayRef<NamedAttribute>", "{{}">:$attributes),
[{{
$_state.addOperands(operands);
$_state.addAttributes(attributes);
$_state.addTypes(resultTensorTypes);
(void)$_state.addRegion();
}]>
{6}
];
let printer = [{{ return ::printNamedStructuredOp(p, *this); }];
let parser = [{{
return ::parseNamedStructuredOp<{0}>(parser, result/*TODO:, captures*/);
}];
let hasFolder = 1;
let hasCanonicalizer = 1;
let extraClassDeclaration = structuredOpsBaseDecls # [{{
// Auto-generated.
ArrayAttr iterator_types();
ArrayAttr indexing_maps();
static void regionBuilder(Block &block, ValueRange captures);
static std::function<void(Block &, ValueRange)> getRegionBuilder() {{
return regionBuilder;
}
// Generic methods.
static unsigned getNumRegionArgs();
std::string getLibraryCallName();
{7}
}];
}
)FMT";
// The iterator_types() method implementation. Parameters:
// {0}: Class name
// {1}: Comma interleaved iterator type names.
static const char structuredOpIteratorTypesFormat[] =
R"FMT(
ArrayAttr {0}::iterator_types() {
return Builder(getContext()).getStrArrayAttr(SmallVector<StringRef>{{ {1} });
}
)FMT";
// Implementations of getCanonicalizationPatterns, fold and getEffects.
// Parameters:
// {0}: Class name
const char structuredOpCanonicalizersAndFoldersFormat[] = R"FMT(
void {0}::getCanonicalizationPatterns(
OwningRewritePatternList &results,
MLIRContext *context) {{
results.insert<EraseDeadLinalgOp>();
results.insert<FoldTensorCastOp>();
}
LogicalResult {0}::fold(ArrayRef<Attribute>,
SmallVectorImpl<OpFoldResult> &) {{
return foldMemRefCast(*this);
}
void {0}::getEffects(SmallVectorImpl<
SideEffects::EffectInstance<MemoryEffects::Effect> >&effects) {{
getGenericEffectsImpl(effects,
getOperation()->getResults(), getInputBuffers(), getOutputBuffers());
}
)FMT";
static LogicalResult generateNamedGenericOpOds(LinalgOpConfig &opConfig,
GenerationContext &genContext) {
if (!genContext.shouldGenerateOds())
return success();
raw_ostream &os = genContext.odss();
std::string interfaceNameList;
std::string attrList;
std::string attrMethods;
std::string attrBuilder;
std::string doc;
if (opConfig.metadata->doc) {
const char *docFmt = R"FMT(
let summary = [{ {0} }];
let description = [{
{1}
}];
)FMT";
StringRef summary, description;
std::tie(summary, description) =
StringRef(*opConfig.metadata->doc).trim().split('\n');
doc = llvm::formatv(docFmt, summary.trim(), description.trim());
}
os << llvm::formatv(structuredOpOdsHeaderFormat, opConfig.metadata->cppOpName,
opConfig.metadata->name, interfaceNameList, doc, attrList,
opConfig.structuredOp->args.size(), attrBuilder,
attrMethods);
return success();
}
static LogicalResult
generateNamedGenericOpDefns(LinalgOpConfig &opConfig,
GenerationContext &genContext) {
if (!genContext.shouldGenerateDefns())
return success();
raw_ostream &os = genContext.defns();
StringRef className = opConfig.metadata->cppOpName;
// Implementation banner.
std::string bannerComment = llvm::formatv("Implementation of {0}", className);
os << llvm::formatv(bannerFormat, bannerComment);
// Reference iterators.
{
std::string iteratorsStr;
llvm::raw_string_ostream ss(iteratorsStr);
llvm::interleaveComma(opConfig.structuredOp->iteratorTypes, ss,
[&](LinalgIteratorTypeDef it) {
switch (it) {
case LinalgIteratorTypeDef::parallel:
ss << "getParallelIteratorTypeName()";
break;
case LinalgIteratorTypeDef::reduction:
ss << "getReductionIteratorTypeName()";
break;
}
});
ss.flush();
os << llvm::formatv(structuredOpIteratorTypesFormat, className,
iteratorsStr);
}
// Static indexing maps.
if (auto &staticMaps =
opConfig.structuredOp->indexingMaps.staticIndexingMaps) {
if (staticMaps->empty())
return emitError(genContext.getLoc()) << "op has no indexing maps";
AffineMap firstMap = staticMaps->front().affineMap();
// Symbol bindings.
{
// For each symbol, generate a declaration for it, either with an
// AffineSymbolExpr or an AffineConstantExpr (if the symbol derives from
// an attribute).
// TODO: Implement attribute constants.
// TODO: Possibly lift into a top-level method.
static const char structuredOpSymbolBindingsFormat[] = R"FMT(
static SmallVector<AffineExpr> getSymbolBindings({0} self) {
MLIRContext *context = self.getContext();
SmallVector<AffineExpr> exprs;
{1}
return exprs;
}
)FMT";
unsigned symbolCount = firstMap.getNumSymbols();
SmallVector<std::string> symbolBindings;
for (unsigned i = 0; i < symbolCount; ++i) {
// TODO: Switch and emit constants for attribute bound symbols.
symbolBindings.push_back(llvm::formatv(
" exprs.push_back(getAffineSymbolExpr({0}, context));", i));
}
std::string symbolBindingsStr;
llvm::raw_string_ostream symbolBindingsSs(symbolBindingsStr);
llvm::interleave(symbolBindings, symbolBindingsSs, "\n");
symbolBindingsSs.flush();
os << llvm::formatv(structuredOpSymbolBindingsFormat, className,
symbolBindingsStr);
}
// Indexing maps.
{
// Parameters:
// {0}: Class name
// {1}: Comma-separated list of dimension variable names.
// {2}: Statements
static const char structuredOpIndexingMapsFormat[] = R"FMT(
ArrayAttr {0}::indexing_maps() {
MLIRContext *context = getContext();
auto symbolBindings = getSymbolBindings(*this);
SmallVector<AffineMap> maps;
{2}
return Builder(context).getAffineMapArrayAttr(maps);
}
)FMT";
unsigned dimCount = firstMap.getNumDims();
// Generate a comma-separated list of dim identifiers to be passed to
// bindDims, ensuring tht AffineExpr identifiers are bound in the right
// order to the proper AffineDimExpr.
// This results in vars in scope like: d0, d1, d2...
SmallVector<unsigned> dimIndices;
for (unsigned i = 0; i < dimCount; ++i)
dimIndices.push_back(i);
std::string dimIdentsStr;
llvm::raw_string_ostream dimIdentsSs(dimIdentsStr);
llvm::interleaveComma(dimIndices, dimIdentsSs,
[&](unsigned i) { dimIdentsSs << "d" << i; });
dimIdentsSs.flush();
// Statements to add and simplify each affine map.
SmallVector<std::string> stmts;
for (auto &indexingMap : *staticMaps) {
// TODO: Assert that dim and symbol count match the first.
stmts.push_back(
llvm::formatv("maps.push_back({0});",
generateCppExpression(indexingMap, "context")));
stmts.push_back(llvm::formatv(
"maps.back() = "
"simplifyAffineMap(maps.back().replaceDimsAndSymbols({{}, "
"symbolBindings, {0}, 0));",
dimCount));
}
// TODO: This needs to be memoized and/or converted to non-parser based
// C++ codegen prior to real use.
os << llvm::formatv(structuredOpIndexingMapsFormat, className,
dimIdentsStr, interleaveToString(stmts, "\n "));
}
} else {
return emitError(genContext.getLoc())
<< "generating code for non static indexing maps not currently "
"supported";
}
// getNumRegionArgs()
{
// Generates a getNumRegionArgs() method. Parameters:
// {0}: Class name
// {1}: Number of region args
static const char structuredOpGetNumRegionArgsFormat[] = R"FMT(
unsigned {0}::getNumRegionArgs() {{ return {1}; }
)FMT";
os << llvm::formatv(structuredOpGetNumRegionArgsFormat, className,
opConfig.structuredOp->args.size());
}
// getLibraryCallName()
{
// Generates a getLibraryCallName method. Parameters:
// {0}: Class name
static const char structuredOpGetLibraryCallFormat[] = R"FMT(
std::string {0}::getLibraryCallName() {{
return generateLibraryCallName(getOperation());
}
)FMT";
os << llvm::formatv(structuredOpGetLibraryCallFormat, className);
}
// regionBuilder()
{
// Generates a regionBuilder method. Parameters.
// {0}: Class name
// {1}: Statements
static const char structuredOpRegionBuilderFormat[] = R"FMT(
void {0}::regionBuilder(Block &block, ValueRange captures) {{
RegionBuilderHelper helper(block);
SmallVector<Value> yields;
{1}
helper.yieldOutputs(yields);
}
)FMT";
auto &args = opConfig.structuredOp->args;
auto &assignments = opConfig.structuredOp->assignments;
size_t generatedAssignmentCount = 0;
int localCounter = 0;
SmallVector<std::string> stmts;
for (LinalgTensorDef &arg : args) {
if (arg.usage != LinalgTensorUsageDef::output &&
arg.usage != LinalgTensorUsageDef::temporary)
continue;
// Find the assignment that correlates with the argument.
ScalarAssign *assignment = findAssignment(arg.name, assignments);
if (!assignment)
return emitError(genContext.getLoc())
<< "no assignment found for output argument " << arg.name;
++generatedAssignmentCount;
// Recursively generate the expression.
std::function<Optional<std::string>(ScalarExpression &)>
generateExpression =
[&](ScalarExpression &expression) -> Optional<std::string> {
if (expression.scalarArg) {
Optional<int> argIndex =
findTensorDefArgIndex(*expression.scalarArg, args);
if (!argIndex) {
emitError(genContext.getLoc())
<< "scalar argument not defined on the op: " << arg.name;
return None;
}
return std::string(
llvm::formatv("block.getArgument({0})", *argIndex));
} else if (expression.scalarApply) {
// Recursively generate operands.
SmallVector<std::string> operandCppValues;
for (ScalarExpression &operand : expression.scalarApply->operands) {
auto operandCppValue = generateExpression(operand);
if (!operandCppValue)
return None;
operandCppValues.push_back(*operandCppValue);
}
std::string cppIdent = llvm::formatv("value{0}", ++localCounter);
stmts.push_back(
llvm::formatv("Value {0} = helper.applyfn__{1}({2});", cppIdent,
expression.scalarApply->fnName,
interleaveToString(operandCppValues, ", ")));
return cppIdent;
} else {
emitError(genContext.getLoc()) << "unknown ScalarExpression type";
return None;
}
};
Optional<std::string> cppValue = generateExpression(assignment->value);
if (!cppValue)
return failure();
stmts.push_back(llvm::formatv("yields.push_back({0});", cppValue));
}
if (generatedAssignmentCount != assignments.size())
return emitError(genContext.getLoc())
<< "mismatched number of assignments vs output arguments";
os << llvm::formatv(structuredOpRegionBuilderFormat, className,
interleaveToString(stmts, "\n "));
}
// Canonicalizers and folders.
os << llvm::formatv(structuredOpCanonicalizersAndFoldersFormat, className);
return success();
}
static LogicalResult generateOp(LinalgOpConfig &opConfig,
GenerationContext &genContext) {
// Switch on op type being generated.
if (opConfig.structuredOp) {
return success(
succeeded(generateNamedGenericOpOds(opConfig, genContext)) &&
succeeded(generateNamedGenericOpDefns(opConfig, genContext)));
} else {
return emitError(genContext.getLoc()) << "unsupported operation type";
}
}
//===----------------------------------------------------------------------===//
// Command line options and main
//===----------------------------------------------------------------------===//
static llvm::cl::opt<std::string>
inputFilename(llvm::cl::Positional, llvm::cl::desc("<input file>"),
llvm::cl::init("-"), llvm::cl::value_desc("YAML filename"));
static llvm::cl::opt<std::string>
outputOdsDeclFilename("o-ods-decl", llvm::cl::desc("ODS output filename"),
llvm::cl::value_desc("filename"), llvm::cl::init(""));
static llvm::cl::opt<std::string>
outputCppImplFilename("o-impl",
llvm::cl::desc("C++ implementation file name"),
llvm::cl::value_desc("filename"), llvm::cl::init(""));
int main(int argc, char **argv) {
llvm::cl::ParseCommandLineOptions(argc, argv, "Linalg ODS Gen from YAML");
// Set up the input file.
std::string errorMessage;
std::unique_ptr<llvm::MemoryBuffer> file =
mlir::openInputFile(inputFilename, &errorMessage);
if (!file) {
llvm::errs() << errorMessage << "\n";
return 1;
}
MLIRContext mlirContext;
LinalgYAMLContext yamlContext{&mlirContext};
std::vector<LinalgOpConfig> opConfigs;
// Parse input.
Input yin(file->getBuffer(), &yamlContext);
yin >> opConfigs;
if (yin.error())
return 1;
// Open output files.
std::unique_ptr<llvm::ToolOutputFile> outputOdsDecl;
if (!outputOdsDeclFilename.empty()) {
outputOdsDecl = openOutputFile(outputOdsDeclFilename, &errorMessage);
if (!outputOdsDecl) {
llvm::errs() << errorMessage << "\n";
return 1;
}
}
std::unique_ptr<llvm::ToolOutputFile> outputCppImpl;
if (!outputCppImplFilename.empty()) {
outputCppImpl = openOutputFile(outputCppImplFilename, &errorMessage);
if (!outputCppImpl) {
llvm::errs() << errorMessage << "\n";
return 1;
}
}
if (!outputOdsDecl && !outputCppImpl) {
llvm::errs() << "error: No output files specified\n";
return 1;
}
// Generate.
GenerationContext genContext(&mlirContext,
outputOdsDecl ? &outputOdsDecl->os() : nullptr,
outputCppImpl ? &outputCppImpl->os() : nullptr);
for (auto &opConfig : opConfigs) {
if (!opConfig.metadata) {
emitError(genContext.getLoc())
<< "missing operation metadata on subsequent op";
return 1;
}
genContext.setLoc(NameLoc::get(
Identifier::get(opConfig.metadata->cppOpName, &mlirContext),
&mlirContext));
if (failed(generateOp(opConfig, genContext))) {
return 1;
}
}
if (outputOdsDecl)
outputOdsDecl->keep();
if (outputCppImpl)
outputCppImpl->keep();
return 0;
}