diff --git a/mlir/include/mlir/Dialect/Tensor/IR/CMakeLists.txt b/mlir/include/mlir/Dialect/Tensor/IR/CMakeLists.txt --- a/mlir/include/mlir/Dialect/Tensor/IR/CMakeLists.txt +++ b/mlir/include/mlir/Dialect/Tensor/IR/CMakeLists.txt @@ -1,2 +1,15 @@ add_mlir_dialect(TensorOps tensor) add_mlir_doc(TensorOps TensorOps Dialects/ -gen-dialect-doc) + +set(LLVM_TARGET_DEFINITIONS TensorAttrDefs.td) +mlir_tablegen(TensorAttrDefs.h.inc -gen-attrdef-decls) +mlir_tablegen(TensorAttrDefs.cpp.inc -gen-attrdef-defs) +add_public_tablegen_target(MLIRTensorAttrDefsIncGen) + +set(LLVM_TARGET_DEFINITIONS TensorInterfaces.td) +mlir_tablegen(TensorEnums.h.inc -gen-enum-decls) +mlir_tablegen(TensorEnums.cpp.inc -gen-enum-defs) +mlir_tablegen(TensorAttrInterfaces.h.inc -gen-attr-interface-decls) +mlir_tablegen(TensorAttrInterfaces.cpp.inc -gen-attr-interface-defs) +add_public_tablegen_target(MLIRTensorInterfacesIncGen) + diff --git a/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h b/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h --- a/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h +++ b/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h @@ -9,6 +9,7 @@ #ifndef MLIR_DIALECT_TENSOR_IR_TENSOR_H_ #define MLIR_DIALECT_TENSOR_IR_TENSOR_H_ +#include "mlir/IR/AffineMap.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Dialect.h" #include "mlir/IR/OpDefinition.h" @@ -23,6 +24,16 @@ #include "mlir/Dialect/Tensor/IR/TensorOpsDialect.h.inc" +//===----------------------------------------------------------------------===// +// Tensor Dialect Enums, Attribute Interfaces, and Attributes +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/Tensor/IR/TensorEnums.h.inc" +// enums before attr +#include "mlir/Dialect/Tensor/IR/TensorAttrInterfaces.h.inc" +#define GET_ATTRDEF_CLASSES +#include "mlir/Dialect/Tensor/IR/TensorAttrDefs.h.inc" + //===----------------------------------------------------------------------===// // Tensor Dialect Operations //===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorAttrDefs.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorAttrDefs.td new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorAttrDefs.td @@ -0,0 +1,34 @@ +//===-- TensorAttrDefs.td - Tensor Attributes Definitions --*- tablegen -*-===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#ifndef TENSOR_ATTRDEFS +#define TENSOR_ATTRDEFS + +include "mlir/Dialect/Tensor/IR/TensorBase.td" +include "mlir/Dialect/Tensor/IR/TensorInterfaces.td" + +// All of the Tensor attributes will extend this class. +class Tensor_Attr traits = []> : AttrDef; + +// Sparse tensor encoding attribute. +def SparseTensorEncodingAttr : Tensor_Attr<"SparseEncoding", + [ DeclareAttrInterfaceMethods ]> { + let mnemonic = "sparse"; + let description = [{ + Attribute to encode sparse tensor information through a dictionary. + The semantics are defined by the `SparseTensorEncoding` interface. + }]; + let parameters = ( + ins + "DictionaryAttr":$dict + ); + +} + +#endif // LLVMIR_ATTRDEFS diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorInterfaces.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorInterfaces.td new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorInterfaces.td @@ -0,0 +1,216 @@ +//===- TensorInterfaces.td - Tensor interfaces -------------*- tablegen -*-===// +// +// 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 +// +//===----------------------------------------------------------------------===// +// +// Defines the interfaces associated with tensor types and attributes. +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_IR_TENSORINTERFACES +#define MLIR_IR_TENSORINTERFACES + +include "mlir/IR/OpBase.td" + +//===----------------------------------------------------------------------===// +// Dimension level types that define sparse tensors: +// Dense - dimension is dense, every entry is stored +// Compressed - dimension is sparse, only nonzeros are stored +// Singleton - dimension contains single coordinate, no siblings +//===----------------------------------------------------------------------===// + +def DimLevelTypeDense : StrEnumAttrCase<"dense">; +def DimLevelTypeCompressed : StrEnumAttrCase<"compressed">; +def DimLevelTypeSingleton : StrEnumAttrCase<"singleton">; + +def DimLevelType : StrEnumAttr< + "DimLevelType", + "dimension level type", + [DimLevelTypeDense, DimLevelTypeCompressed, DimLevelTypeSingleton]> { + let cppNamespace = "::mlir"; +} + +//===----------------------------------------------------------------------===// +// Attribute interface +//===----------------------------------------------------------------------===// + +def SparseTensorEncoding : AttrInterface<"SparseTensorEncoding"> { + let cppNamespace = "::mlir"; + let description = [{ + This defines an interface that assigns sparse semantics to a + tensor.sparse<> encoding that is associated with a tensor type. + A "TACO"-style (see tensor-compiler.org) way of informing the + compiler of the sparsity of tensors is supported by means of the + methods getDimLevelType(), getDimOrdering(), getPointerType(), and + getIndexType(). The semantics of these methods are described below. + Clients should first assert that isValid() succeeds for + the tensor size dimension. + + The encoding is eventually used by a `sparse compiler` pass + to generate sparse code fully automatically for all tensor + expressions that involve tensors with a sparse encoding. + Compiler passes that run before this sparse compiler pass + need to be aware of the semantics of tensor types with such + an encoding. + }]; + let methods = [ + InterfaceMethod< + /*desc=*/[{ + Verifies the encoding is valid for a tensor with the given dimension + size. Reports an error at the given location otherwise. + }], + /*retTy=*/"::mlir::LogicalResult", + /*methodName=*/"isValid", + /*args=*/(ins "unsigned":$size, "Location":$loc), + /*methodBody=*/"", + /*defaultImplementation=*/[{ + for (const NamedAttribute &attr : $_attr.getDict()) { + if (attr.first == getSparseDimLevelTypeAttrName()) { + // + // Dimension level type verification. + // + auto arrayAttr = attr.second.dyn_cast(); + if (!arrayAttr || size != static_cast(arrayAttr.size())) + return emitError(loc) << "expected an array of size " << size << " for dimension level types"; + for (unsigned i = 0; i < size; i++) { + auto strAttr = arrayAttr[i].dyn_cast(); + if (!strAttr) + return emitError(loc) << "expected string value in dimension level types"; + if (!symbolizeDimLevelType(strAttr.getValue()).hasValue()) + return emitError(loc) << "unexpected dimension level type: " << strAttr; + } + } else if (attr.first == getSparseDimOrderingAttrName()) { + // + // Dimension order verification. + // + auto affineAttr = attr.second.dyn_cast(); + if (!affineAttr) + return emitError(loc) << "expected an affine map for dimension ordering"; + AffineMap map = affineAttr.getValue(); + if (size != map.getNumResults() || !map.isPermutation()) + return emitError(loc) << "expected a permutation affine map of size " << size << " for dimension ordering"; + } else if (attr.first == getSparsePointerBitWidthAttrName() || + attr.first == getSparseIndexBitWidthAttrName()) { + // + // Pointer or index bitwidth verification. + // + auto intAttr = attr.second.dyn_cast(); + if (!intAttr) + return emitError(loc) << "expected an integral bitwidth"; + switch (intAttr.getInt()) { + case 0: case 8: case 16: case 32: case 64: + continue; + default: + return emitError(loc) << "unexpected bitwidth: " << intAttr.getInt(); + } + } else { + return emitError(loc) << "unexpected key: " << attr.first.str(); + } + } + return success(); + }] + >, + InterfaceMethod< + /*desc=*/[{ + Returns the dimension level type in the given dimension `dim` + of this tensor type. The choices, defined by the `DimLevelType` + enum, are `dense` (the dimension should be stored in its entirety), + `compressed` (only non-zero regions or elements should be stored), + or `singleton` (no sibling elements for parent). + }], + /*retTy=*/"::mlir::DimLevelType", + /*methodName=*/"getDimLevelType", + /*args=*/(ins "unsigned":$dim), + /*methodBody=*/"", + /*defaultImplementation=*/[{ + if (auto value = $_attr.getDict().get( + SparseTensorEncoding::getSparseDimLevelTypeAttrName())) { + auto optionalValue = symbolizeDimLevelType( + value.template cast()[dim].template cast().getValue()); + if (optionalValue.hasValue()) + return optionalValue.getValue(); + } + return DimLevelType::dense; + }] + >, + InterfaceMethod< + /*desc=*/[{ + Returns the dimension order of this tensor type as an AffineMap. + Unlike dense storage, most sparse storage schemes do not provide + fast random access. This affine map specifies the order of + dimensions that should be support by the sparse storage scheme + (e.g. (i,j) -> (i,j) requests 2-d row-wise and (i,j) -> (j,i) + requests 2-d column-wise storage). + }], + /*retTy=*/"::mlir::AffineMap", + /*methodName=*/"getDimOrdering", + /*args=*/(ins), + /*methodBody=*/"", + /*defaultImplementation=*/[{ + if (auto value = $_attr.getDict().get( + SparseTensorEncoding::getSparseDimOrderingAttrName())) + return value.template cast().getValue(); + return {}; + }] + >, + InterfaceMethod< + /*desc=*/[{ + Returns the required bit width for pointer storage. A narrow width + reduces the memory footprint of overhead storage, as long as the + width suffices to define the total required range (viz. the maximum + number of stored entries over all indirection dimensions). The choices + are `8`, `16`, `32`, `64`, or the default `0`, which selects the + native index width for pointers. + }], + /*retTy=*/"unsigned", + /*methodName=*/"getPointerBitWidth", + /*args=*/(ins), + /*methodBody=*/"", + /*defaultImplementation=*/[{ + if (auto value = $_attr.getDict().get( + SparseTensorEncoding::getSparsePointerBitWidthAttrName())) + return value.template cast().getInt(); + return 0; + }] + >, + InterfaceMethod< + /*desc=*/[{ + Returns the required bit width for index storage. A narrow width + reduces the memory footprint of overhead storage, as long as the + width suffices to define the total required range (viz. the maximum + value of each tensor index over all dimensions). The choices are `8`, + `16`, `32`, `64`, or the default `0`, which selects the native index + width for indices. + }], + /*retTy=*/"unsigned", + /*methodName=*/"getIndexBitWidth", + /*args=*/(ins), + /*methodBody=*/"", + /*defaultImplementation=*/[{ + if (auto value = $_attr.getDict().get( + SparseTensorEncoding::getSparseIndexBitWidthAttrName())) + return value.template cast().getInt(); + return 0; + }] + > + ]; + let extraClassDeclaration = [{ + constexpr static StringRef getSparseDimLevelTypeAttrName() { + return "sparseDimLevelType"; + } + constexpr static StringRef getSparseDimOrderingAttrName() { + return "sparseDimOrdering"; + } + constexpr static StringRef getSparsePointerBitWidthAttrName() { + return "sparsePointerBitWidth"; + } + constexpr static StringRef getSparseIndexBitWidthAttrName() { + return "sparseIndexBitWidth"; + } + }]; +} + +#endif // MLIR_IR_TENSORINTERFACES diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td --- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td +++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td @@ -10,6 +10,7 @@ #define TENSOR_OPS include "mlir/Dialect/Tensor/IR/TensorBase.td" +include "mlir/Dialect/Tensor/IR/TensorAttrDefs.td" include "mlir/Interfaces/CastInterfaces.td" include "mlir/Interfaces/ControlFlowInterfaces.td" include "mlir/Interfaces/SideEffectInterfaces.td" diff --git a/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt b/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt --- a/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt +++ b/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt @@ -7,6 +7,8 @@ DEPENDS MLIRTensorOpsIncGen + MLIRTensorAttrDefsIncGen + MLIRTensorInterfacesIncGen LINK_COMPONENTS Core diff --git a/mlir/lib/Dialect/Tensor/IR/TensorDialect.cpp b/mlir/lib/Dialect/Tensor/IR/TensorDialect.cpp --- a/mlir/lib/Dialect/Tensor/IR/TensorDialect.cpp +++ b/mlir/lib/Dialect/Tensor/IR/TensorDialect.cpp @@ -7,11 +7,38 @@ //===----------------------------------------------------------------------===// #include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/DialectImplementation.h" #include "mlir/Transforms/InliningUtils.h" +#include "llvm/ADT/TypeSwitch.h" using namespace mlir; using namespace mlir::tensor; +//===----------------------------------------------------------------------===// +// TableGen'd Enums, Attribute Interfaces, and Attributes Methods +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/Tensor/IR/TensorAttrInterfaces.cpp.inc" +#include "mlir/Dialect/Tensor/IR/TensorEnums.cpp.inc" +#define GET_ATTRDEF_CLASSES +#include "mlir/Dialect/Tensor/IR/TensorAttrDefs.cpp.inc" + +Attribute SparseEncodingAttr::parse(MLIRContext *context, + DialectAsmParser &parser, Type type) { + if (failed(parser.parseLess())) + return {}; + DictionaryAttr dict; + if (failed(parser.parseAttribute(dict))) + return {}; + if (failed(parser.parseGreater())) + return {}; + return SparseEncodingAttr::get(context, dict); +} + +void SparseEncodingAttr::print(DialectAsmPrinter &printer) const { + printer << "sparse<" << getDict() << ">"; +} + //===----------------------------------------------------------------------===// // TensorDialect Dialect Interfaces //===----------------------------------------------------------------------===// @@ -30,10 +57,38 @@ }; } // end anonymous namespace +//===----------------------------------------------------------------------===// +// TensorDialect Methods +//===----------------------------------------------------------------------===// + void TensorDialect::initialize() { + addAttributes< +#define GET_ATTRDEF_LIST +#include "mlir/Dialect/Tensor/IR/TensorAttrDefs.cpp.inc" + >(); addOperations< #define GET_OP_LIST #include "mlir/Dialect/Tensor/IR/TensorOps.cpp.inc" >(); addInterfaces(); } + +Attribute TensorDialect::parseAttribute(DialectAsmParser &parser, + Type type) const { + StringRef attrTag; + if (failed(parser.parseKeyword(&attrTag))) + return Attribute(); + Attribute attr; + auto parseResult = + generatedAttributeParser(getContext(), parser, attrTag, type, attr); + if (parseResult.hasValue()) + return attr; + parser.emitError(parser.getNameLoc(), "unknown tensor attribute"); + return Attribute(); +} + +void TensorDialect::printAttribute(::mlir::Attribute attr, + ::mlir::DialectAsmPrinter &printer) const { + if (succeeded(generatedAttributePrinter(attr, printer))) + return; +} diff --git a/mlir/lib/IR/BuiltinTypes.cpp b/mlir/lib/IR/BuiltinTypes.cpp --- a/mlir/lib/IR/BuiltinTypes.cpp +++ b/mlir/lib/IR/BuiltinTypes.cpp @@ -446,7 +446,9 @@ for (int64_t s : shape) if (s < -1) return emitError() << "invalid tensor dimension size"; - // TODO: verify contents of encoding attribute. + // TODO + // if (failed(isValidEncoding(encoding, shape.size(), emitError))) + // return failure(); return checkTensorElementType(emitError, elementType); } diff --git a/mlir/lib/Parser/TypeParser.cpp b/mlir/lib/Parser/TypeParser.cpp --- a/mlir/lib/Parser/TypeParser.cpp +++ b/mlir/lib/Parser/TypeParser.cpp @@ -412,8 +412,12 @@ // Parse an optional encoding attribute. Attribute encoding; - if (consumeIf(Token::comma)) + if (consumeIf(Token::comma)) { encoding = parseAttribute(); + // if (failed(isValidEncoding(encoding, dimensions.size(), + // [&] { return emitError(); }))) + // return nullptr; + } if (!elementType || parseToken(Token::greater, "expected '>' in tensor type")) return nullptr; diff --git a/mlir/test/Dialect/Tensor/invalid_sparse_tensor.mlir b/mlir/test/Dialect/Tensor/invalid_sparse_tensor.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/Tensor/invalid_sparse_tensor.mlir @@ -0,0 +1,58 @@ +// RUN: mlir-opt <%s -test-sparse-encoding -split-input-file -verify-diagnostics + +// ----- + +#a = #tensor.sparse<{sparseDimLevelType = [1,2]}> +func private @tensor_size_mismatch(%arg0: tensor<8xi32, #a>) -> () { // expected-error {{expected an array of size 1 for dimension level types}} + return +} + +// ----- + +#a = #tensor.sparse<{sparseDimLevelType = [1]}> +func private @tensor_type_mismatch(%arg0: tensor<8xi32, #a>) -> () { // expected-error {{expected string value in dimension level types}} + return +} + +// ----- + +#a = #tensor.sparse<{sparseDimLevelType = ["strange"]}> +func private @tensor_value_mismatch(%arg0: tensor<8xi32, #a>) -> () { // expected-error {{unexpected dimension level type: "strange"}} + return +} + +// ----- + +#a = #tensor.sparse<{sparseDimOrdering = "wrong"}> +func private @tensor_order_mismatch(%arg0: tensor<8xi32, #a>) -> () { // expected-error {{expected an affine map for dimension ordering}} + return +} + +// ----- + +#a = #tensor.sparse<{sparseDimOrdering = affine_map<(i,j) -> (i,i)>}> +func private @tensor_no_permutation(%arg0: tensor<16x32xf32, #a>) -> () { // expected-error {{expected a permutation affine map of size 2 for dimension ordering}} + return +} + +// ----- + +#a = #tensor.sparse<{sparseIndexBitWidth = "not really"}> +func private @tensor_no_int(%arg0: tensor<16x32xf32, #a>) -> () { // expected-error {{expected an integral bitwidth}} + return +} + +// ----- + +#a = #tensor.sparse<{sparseIndexBitWidth = 128}> +func private @tensor_invalid_int(%arg0: tensor<16x32xf32, #a>) -> () { // expected-error {{unexpected bitwidth: 128}} + return +} + +// ----- + +#a = #tensor.sparse<{key = 1}> +func private @tensor_invalid_key(%arg0: tensor<16x32xf32, #a>) -> () { // expected-error {{unexpected key: key}} + return +} + diff --git a/mlir/test/lib/Transforms/TestSparsification.cpp b/mlir/test/lib/Transforms/TestSparsification.cpp --- a/mlir/test/lib/Transforms/TestSparsification.cpp +++ b/mlir/test/lib/Transforms/TestSparsification.cpp @@ -127,6 +127,31 @@ } }; +// TODO: define a proper tensortype <-> encoding verification hook +// and replace this testing code by that verification +struct TestSparseEncoding + : public PassWrapper { + void runOnFunction() override { + Location loc = getFunction().getLoc(); + for (Block &block : getFunction().getBlocks()) { + for (Value v : block.getArguments()) + test(v, loc); + for (Operation &op : block) + for (Value v : op.getOperands()) + test(v, loc); + } + } + void test(Value v, Location loc) { + if (auto rtp = v.getType().dyn_cast()) { + unsigned size = rtp.getShape().size(); + if (auto intf = + rtp.getEncoding().dyn_cast_or_null()) + if (failed(intf.isValid(size, loc))) + this->signalPassFailure(); + } + } +}; + } // end anonymous namespace namespace mlir { @@ -135,6 +160,8 @@ void registerTestSparsification() { PassRegistration sparsificationPass( "test-sparsification", "Test automatic generation of sparse tensor code"); + PassRegistration sparseEncodingPass( + "test-sparse-encoding", "Test sparse encoding on tensor types"); } } // namespace test