diff --git a/mlir/include/mlir/Dialect/CMakeLists.txt b/mlir/include/mlir/Dialect/CMakeLists.txt --- a/mlir/include/mlir/Dialect/CMakeLists.txt +++ b/mlir/include/mlir/Dialect/CMakeLists.txt @@ -14,5 +14,6 @@ add_subdirectory(Shape) add_subdirectory(SPIRV) add_subdirectory(StandardOps) +add_subdirectory(Tensor) add_subdirectory(Tosa) add_subdirectory(Vector) diff --git a/mlir/include/mlir/Dialect/Linalg/EDSC/FoldedIntrinsics.h b/mlir/include/mlir/Dialect/Linalg/EDSC/FoldedIntrinsics.h --- a/mlir/include/mlir/Dialect/Linalg/EDSC/FoldedIntrinsics.h +++ b/mlir/include/mlir/Dialect/Linalg/EDSC/FoldedIntrinsics.h @@ -10,6 +10,7 @@ #include "mlir/Dialect/Linalg/EDSC/Builders.h" #include "mlir/Dialect/Linalg/EDSC/Intrinsics.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Transforms/FoldUtils.h" @@ -46,7 +47,6 @@ using folded_std_constant_index = FoldedValueBuilder; using folded_std_constant_int = FoldedValueBuilder; using folded_std_dim = FoldedValueBuilder; -using folded_std_extract_element = FoldedValueBuilder; using folded_std_index_cast = FoldedValueBuilder; using folded_std_muli = FoldedValueBuilder; using folded_std_mulf = FoldedValueBuilder; @@ -60,6 +60,7 @@ using folded_std_view = FoldedValueBuilder; using folded_std_zero_extendi = FoldedValueBuilder; using folded_std_sign_extendi = FoldedValueBuilder; +using folded_tensor_extract = FoldedValueBuilder; } // namespace intrinsics } // namespace edsc } // namespace mlir diff --git a/mlir/include/mlir/Dialect/StandardOps/EDSC/Intrinsics.h b/mlir/include/mlir/Dialect/StandardOps/EDSC/Intrinsics.h --- a/mlir/include/mlir/Dialect/StandardOps/EDSC/Intrinsics.h +++ b/mlir/include/mlir/Dialect/StandardOps/EDSC/Intrinsics.h @@ -9,6 +9,7 @@ #define MLIR_DIALECT_STANDARDOPS_EDSC_INTRINSICS_H_ #include "mlir/Dialect/StandardOps/EDSC/Builders.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" namespace mlir { namespace edsc { @@ -28,7 +29,6 @@ using std_divis = ValueBuilder; using std_diviu = ValueBuilder; using std_dim = ValueBuilder; -using std_extract_element = ValueBuilder; using std_fpext = ValueBuilder; using std_fptrunc = ValueBuilder; using std_im = ValueBuilder; @@ -52,6 +52,7 @@ using std_view = ValueBuilder; using std_zero_extendi = ValueBuilder; using std_sign_extendi = ValueBuilder; +using tensor_extract = ValueBuilder; /// Branches into `block` with `operands`. BranchOp std_br(Block *block, ValueRange operands); diff --git a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td --- a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td +++ b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td @@ -1669,59 +1669,6 @@ let summary = "base-2 exponential of the specified value"; } -//===----------------------------------------------------------------------===// -// ExtractElementOp -//===----------------------------------------------------------------------===// - -def ExtractElementOp : Std_Op<"extract_element", - [NoSideEffect, - TypesMatchWith<"result type matches element type of aggregate", - "aggregate", "result", - "$_self.cast().getElementType()">]> { - let summary = "element extract operation"; - let description = [{ - The `extract_element` op reads a tensor or vector and returns one element - from it specified by an index list. The output of the 'extract_element' is a - new value with the same type as the elements of the tensor or vector. The - arity of indices matches the rank of the accessed value (i.e., if a tensor - is of rank 3, then 3 indices are required for the extract. The indices - should all be of `index` type. - - Example: - - ```mlir - %3 = extract_element %v[%1, %2] : vector<4x4xi32> - %4 = extract_element %t[%1, %2] : tensor<4x4xi32> - %5 = extract_element %ut[%1, %2] : tensor<*xi32> - ``` - }]; - - let arguments = (ins AnyTypeOf<[AnyVector, AnyTensor]>:$aggregate, - Variadic:$indices); - let results = (outs AnyType:$result); - - let builders = [ - OpBuilderDAG<(ins "Value":$aggregate, CArg<"ValueRange", "{}">:$indices), [{ - auto resType = aggregate.getType().cast() - .getElementType(); - build($_builder, $_state, resType, aggregate, indices); - }]>]; - - let extraClassDeclaration = [{ - Value getAggregate() { return getOperand(0); } - - operand_range getIndices() { - return {operand_begin() + 1, operand_end()}; - } - }]; - - let hasFolder = 1; - - let assemblyFormat = [{ - $aggregate `[` $indices `]` attr-dict `:` type($aggregate) - }]; -} - //===----------------------------------------------------------------------===// // TensorFromElementsOp //===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/Dialect/Tensor/CMakeLists.txt b/mlir/include/mlir/Dialect/Tensor/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/Tensor/CMakeLists.txt @@ -0,0 +1,2 @@ +add_subdirectory(IR) +add_subdirectory(Transforms) diff --git a/mlir/include/mlir/Dialect/Tensor/IR/CMakeLists.txt b/mlir/include/mlir/Dialect/Tensor/IR/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/Tensor/IR/CMakeLists.txt @@ -0,0 +1,2 @@ +add_mlir_dialect(TensorOps tensor) +add_mlir_doc(TensorOps -gen-dialect-doc TensorOps Dialects/) diff --git a/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h b/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h @@ -0,0 +1,31 @@ +//===- Tensor.h - Tensor dialect --------------------------------*- C++ -*-===// +// +// 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 MLIR_DIALECT_TENSOR_IR_TENSOR_H_ +#define MLIR_DIALECT_TENSOR_IR_TENSOR_H_ + +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Dialect.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/OpImplementation.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" + +//===----------------------------------------------------------------------===// +// Tensor Dialect +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/Tensor/IR/TensorOpsDialect.h.inc" + +//===----------------------------------------------------------------------===// +// Tensor Dialect Operations +//===----------------------------------------------------------------------===// + +#define GET_OP_CLASSES +#include "mlir/Dialect/Tensor/IR/TensorOps.h.inc" + +#endif // MLIR_DIALECT_TENSOR_IR_TENSOR_H_ diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorBase.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorBase.td new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorBase.td @@ -0,0 +1,48 @@ +//===- TensorBase.td - Base definitions for tensor dialect -*- 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_BASE +#define TENSOR_BASE + +include "mlir/IR/OpBase.td" + +def Tensor_Dialect : Dialect { + let name = "tensor"; + let cppNamespace = "::mlir::tensor"; + let description = [{ + The `tensor` dialect is intended to hold core tensor creation and + manipulation ops, which are not strongly associated with any particular + other dialect or domain abstraction. The primary smoke test of this is ops + that make sense for any tensor element type. + + We leave it to other dialects to hold the vast swath of possible + computations one might want to do on a tensor. + + The `tensor` type is (for better or for worse) used to represent all kinds + of things, and supports an open-ended set of element types. Examples: + + - representing large, dense aggregations of primitive types, suitable for + high-performance numerical computing. + - representing shapes in the `shape` dialect, which consist of small + 1D tensors of `index` data type. + - representing aggregations of strings or “variant” types. + - representing large, sparse aggregations of primitive types, suitable + for high-performance numerical computing. + + Thus, for the `tensor` dialect, we prefer for now to constrain the + scope as much as possible. The expectation is that at some point + in the future, the `tensor` dialect’s scope may be broadened through a + careful discussion of the tradeoffs. + + The `tensor` type is actually a builtin type (it lives in the builtin + dialect), and does not live in this dialect. + + }]; +} + +#endif // TENSOR_BASE diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td @@ -0,0 +1,62 @@ +//===- TensorOps.td - Tensor op 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_OPS +#define TENSOR_OPS + +include "mlir/Dialect/Tensor/IR/TensorBase.td" +include "mlir/Interfaces/SideEffectInterfaces.td" + +class Tensor_Op traits = []> + : Op { + let printer = [{ return ::print(p, *this); }]; + let verifier = [{ return ::verify(*this); }]; + let parser = [{ return ::parse$cppClass(parser, result); }]; +} + +//===----------------------------------------------------------------------===// +// ExtractOp +//===----------------------------------------------------------------------===// + +def Tensor_ExtractOp : Tensor_Op<"extract", + [NoSideEffect, + TypesMatchWith<"result type matches element type of tensor", + "tensor", "result", + "$_self.cast().getElementType()">]> { + let summary = "element extraction operation"; + let description = [{ + The `tensor.extract` op reads a tensor and returns one + element from it specified by an index list. The output of the op is a + new value with the same type as the elements of the tensor. The + arity of indices must match the rank of the accessed value (i.e., if a + tensor is of rank 3, then 3 indices are required for the extract. The + indices should all be of `index` type. + + Example: + + ```mlir + %4 = tensor.extract %t[%1, %2] : tensor<4x4xi32> + %5 = tensor.extract %rt[%1, %2] : tensor + %6 = tensor.extract %ut[%1, %2] : tensor<*xi32> + ``` + }]; + + let arguments = (ins AnyTensor:$tensor, Variadic:$indices); + let results = (outs AnyType:$result); + let assemblyFormat = "$tensor `[` $indices `]` attr-dict `:` type($tensor)"; + + let builders = [ + OpBuilderDAG<(ins "Value":$tensor, CArg<"ValueRange", "{}">:$indices), [{ + auto resType = tensor.getType().cast().getElementType(); + build($_builder, $_state, resType, tensor, indices); + }]>]; + + let hasFolder = 1; +} + +#endif // TENSOR_OPS diff --git a/mlir/include/mlir/Dialect/Tensor/Transforms/CMakeLists.txt b/mlir/include/mlir/Dialect/Tensor/Transforms/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/Tensor/Transforms/CMakeLists.txt @@ -0,0 +1,5 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls -name Tensor) +add_public_tablegen_target(MLIRTensorTransformsIncGen) + +add_mlir_doc(Passes -gen-pass-doc TensorPasses ./) diff --git a/mlir/include/mlir/Dialect/Tensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/Tensor/Transforms/Passes.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/Tensor/Transforms/Passes.h @@ -0,0 +1,38 @@ +//===- Passes.h - Pass Entrypoints ------------------------------*- C++ -*-===// +// +// 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 MLIR_DIALECT_TENSOR_TRANSFORMS_PASSES_H_ +#define MLIR_DIALECT_TENSOR_TRANSFORMS_PASSES_H_ + +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/Bufferize.h" + +namespace mlir { + +class OwningRewritePatternList; + +void populateTensorBufferizePatterns(MLIRContext *context, + BufferizeTypeConverter &typeConverter, + OwningRewritePatternList &patterns); + +/// Creates an instance of `tensor` dialect bufferization pass. +std::unique_ptr createTensorBufferizePass(); + +//===----------------------------------------------------------------------===// +// Registration +//===----------------------------------------------------------------------===// + +namespace tensor { +/// Generate the code for registering passes. +#define GEN_PASS_REGISTRATION +#include "mlir/Dialect/Tensor/Transforms/Passes.h.inc" +} // namespace tensor + +} // end namespace mlir + +#endif // MLIR_DIALECT_TENSOR_TRANSFORMS_PASSES_H_ diff --git a/mlir/include/mlir/Dialect/Tensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/Tensor/Transforms/Passes.td new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/Tensor/Transforms/Passes.td @@ -0,0 +1,19 @@ +//===-- Passes.td - pass definition file -------------------*- 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 MLIR_DIALECT_TENSOR_TRANSFORMS_PASSES +#define MLIR_DIALECT_TENSOR_TRANSFORMS_PASSES + +include "mlir/Pass/PassBase.td" + +def TensorBufferize : FunctionPass<"tensor-bufferize"> { + let summary = "Bufferize the `tensor` dialect"; + let constructor = "mlir::createTensorBufferizePass()"; +} + +#endif // MLIR_DIALECT_TENSOR_TRANSFORMS_PASSES diff --git a/mlir/include/mlir/InitAllDialects.h b/mlir/include/mlir/InitAllDialects.h --- a/mlir/include/mlir/InitAllDialects.h +++ b/mlir/include/mlir/InitAllDialects.h @@ -35,6 +35,7 @@ #include "mlir/Dialect/SPIRV/SPIRVDialect.h" #include "mlir/Dialect/Shape/IR/Shape.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Tosa/IR/TosaOps.h" #include "mlir/Dialect/Vector/VectorOps.h" #include "mlir/IR/Dialect.h" @@ -66,6 +67,7 @@ ROCDL::ROCDLDialect, SDBMDialect, shape::ShapeDialect, + tensor::TensorDialect, tosa::TosaDialect>(); // clang-format on } diff --git a/mlir/include/mlir/InitAllPasses.h b/mlir/include/mlir/InitAllPasses.h --- a/mlir/include/mlir/InitAllPasses.h +++ b/mlir/include/mlir/InitAllPasses.h @@ -25,6 +25,7 @@ #include "mlir/Dialect/SPIRV/Passes.h" #include "mlir/Dialect/Shape/Transforms/Passes.h" #include "mlir/Dialect/StandardOps/Transforms/Passes.h" +#include "mlir/Dialect/Tensor/Transforms/Passes.h" #include "mlir/Dialect/Tosa/Transforms/Passes.h" #include "mlir/Transforms/Passes.h" @@ -57,6 +58,7 @@ registerShapePasses(); spirv::registerSPIRVPasses(); registerStandardPasses(); + tensor::registerTensorPasses(); tosa::registerTosaOptPasses(); } diff --git a/mlir/lib/Conversion/ShapeToStandard/ConvertShapeConstraints.cpp b/mlir/lib/Conversion/ShapeToStandard/ConvertShapeConstraints.cpp --- a/mlir/lib/Conversion/ShapeToStandard/ConvertShapeConstraints.cpp +++ b/mlir/lib/Conversion/ShapeToStandard/ConvertShapeConstraints.cpp @@ -12,6 +12,7 @@ #include "mlir/Dialect/SCF/SCF.h" #include "mlir/Dialect/Shape/IR/Shape.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/PatternMatch.h" #include "mlir/Pass/Pass.h" #include "mlir/Pass/PassRegistry.h" @@ -64,10 +65,10 @@ rewriter.create( loc, rankDiff, greaterRank, one, llvm::None, [&](OpBuilder &b, Location loc, Value iv, ValueRange) { - Value greaterRankOperandExtent = b.create( + Value greaterRankOperandExtent = b.create( loc, greaterRankOperand, ValueRange{iv}); Value ivShifted = b.create(loc, indexTy, iv, rankDiff); - Value lesserRankOperandExtent = b.create( + Value lesserRankOperandExtent = b.create( loc, lesserRankOperand, ValueRange{ivShifted}); Value greaterRankOperandExtentIsOne = b.create( diff --git a/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp b/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp --- a/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp +++ b/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp @@ -12,6 +12,7 @@ #include "mlir/Dialect/SCF/SCF.h" #include "mlir/Dialect/Shape/IR/Shape.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/BlockAndValueMapping.h" #include "mlir/Transforms/DialectConversion.h" @@ -118,12 +119,12 @@ Value outputDimension = args[0]; Value isUnchallengedDimension = b.create( loc, CmpIPredicate::ult, outputDimension, rankDiff); - Value greaterRankOperandExtent = b.create( + Value greaterRankOperandExtent = b.create( loc, greaterRankOperand, outputDimension); // The initial dimensions of the greater-rank operand are unchallenged, // so we can take them as-is. Otherwise, we need to do a comparison. // We need an actual branch here (instead of a select) because the - // lesser-rank operand might be rank 0, so any extract_element would be + // lesser-rank operand might be rank 0, so any tensor.extract would be // invalid. auto ifOp = b.create( loc, TypeRange{indexTy}, isUnchallengedDimension, @@ -140,7 +141,7 @@ // dimensions of zero extent. Value lesserRankOperandDimension = b.create(loc, indexTy, outputDimension, rankDiff); - Value lesserRankOperandExtent = b.create( + Value lesserRankOperandExtent = b.create( loc, lesserRankOperand, ValueRange{lesserRankOperandDimension}); Value greaterRankOperandExtentIsOne = b.create( @@ -262,12 +263,12 @@ auto reduceResult = rewriter.create( loc, rankDiff, greaterRank, one, ValueRange{init}, [&](OpBuilder &b, Location loc, Value iv, ValueRange iterArgs) { - Value greaterRankOperandExtent = - b.create(loc, greaterRankOperand, ValueRange{iv}); + Value greaterRankOperandExtent = b.create( + loc, greaterRankOperand, ValueRange{iv}); Value greaterRankOperandExtentIsOne = b.create( loc, CmpIPredicate::eq, greaterRankOperandExtent, one); Value ivShifted = b.create(loc, indexTy, iv, rankDiff); - Value lesserRankOperandExtent = b.create( + Value lesserRankOperandExtent = b.create( loc, lesserRankOperand, ValueRange{ivShifted}); Value lesserRankOperandExtentIsOne = b.create( loc, CmpIPredicate::eq, lesserRankOperandExtent, one); @@ -316,9 +317,9 @@ } } - rewriter.replaceOpWithNewOp(op, rewriter.getIndexType(), - transformed.shape(), - ValueRange{transformed.dim()}); + rewriter.replaceOpWithNewOp(op, rewriter.getIndexType(), + transformed.shape(), + ValueRange{transformed.dim()}); return success(); } @@ -375,7 +376,8 @@ auto loop = rewriter.create( loc, zero, rank, one, op.initVals(), [&](OpBuilder &b, Location loc, Value iv, ValueRange args) { - Value extent = b.create(loc, transformed.shape(), iv); + Value extent = + b.create(loc, transformed.shape(), iv); SmallVector mappedValues{iv, extent}; mappedValues.append(args.begin(), args.end()); @@ -415,8 +417,8 @@ /// %c1 = constant 1 : index /// %true = constant true /// %4 = scf.for %arg2 = %c0 to %0 step %c1 iter_args(%arg3 = %true) -> (i1) { -/// %5 = extract_element %arg0[%arg2] : tensor -/// %6 = extract_element %arg1[%arg2] : tensor +/// %5 = tensor.extract %arg0[%arg2] : tensor +/// %6 = tensor.extract %arg1[%arg2] : tensor /// %7 = cmpi "eq", %5, %6 : index /// %8 = and %arg3, %7 : i1 /// scf.yield %8 : i1 @@ -465,9 +467,9 @@ [&](OpBuilder &b, Location nestedLoc, Value iv, ValueRange args) { Value conj = args[0]; Value lhsExtent = - b.create(loc, transformed.lhs(), iv); + b.create(loc, transformed.lhs(), iv); Value rhsExtent = - b.create(loc, transformed.rhs(), iv); + b.create(loc, transformed.rhs(), iv); Value eqExtent = b.create(loc, CmpIPredicate::eq, lhsExtent, rhsExtent); Value conjNext = b.create(loc, conj, eqExtent); @@ -584,7 +586,8 @@ // Setup target legality. MLIRContext &ctx = getContext(); ConversionTarget target(ctx); - target.addLegalDialect(); + target + .addLegalDialect(); target.addLegalOp(); // Setup conversion patterns. diff --git a/mlir/lib/Dialect/CMakeLists.txt b/mlir/lib/Dialect/CMakeLists.txt --- a/mlir/lib/Dialect/CMakeLists.txt +++ b/mlir/lib/Dialect/CMakeLists.txt @@ -15,6 +15,7 @@ add_subdirectory(Shape) add_subdirectory(SPIRV) add_subdirectory(StandardOps) +add_subdirectory(Tensor) add_subdirectory(Tosa) add_subdirectory(Vector) diff --git a/mlir/lib/Dialect/StandardOps/CMakeLists.txt b/mlir/lib/Dialect/StandardOps/CMakeLists.txt --- a/mlir/lib/Dialect/StandardOps/CMakeLists.txt +++ b/mlir/lib/Dialect/StandardOps/CMakeLists.txt @@ -15,6 +15,7 @@ MLIREDSC MLIRIR MLIRSideEffectInterfaces + MLIRTensor MLIRVectorInterfaces MLIRViewLikeInterface ) diff --git a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp --- a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp +++ b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp @@ -9,6 +9,7 @@ #include "mlir/Dialect/StandardOps/IR/Ops.h" #include "mlir/Dialect/CommonFolders.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/AffineExpr.h" #include "mlir/IR/AffineMap.h" #include "mlir/IR/BlockAndValueMapping.h" @@ -153,6 +154,7 @@ } void StandardOpsDialect::initialize() { + getContext()->loadDialect(); addOperations /// yield %1 : index /// } : tensor -/// %extracted_element = extract_element %tensor[%c0] : tensor +/// %extracted_element = tensor.extract %tensor[%c0] : tensor /// /// to just with %arg0 replaced by %c0. We only do this if the /// dynamic_tensor_from_elements operation has no side-effects. -struct ExtractElementFromDynamicTensorFromElements - : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; +struct ExtractFromDynamicTensorFromElements + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(ExtractElementOp extract, + LogicalResult matchAndRewrite(tensor::ExtractOp extract, PatternRewriter &rewriter) const final { auto tensorFromElements = - extract.aggregate().getDefiningOp(); + extract.tensor().getDefiningOp(); if (!tensorFromElements || !wouldOpBeTriviallyDead(tensorFromElements)) return failure(); @@ -1894,23 +1896,22 @@ /// Canonicalizes the pattern of the form /// /// %val = tensor_cast %source : : tensor to tensor<2xi32> -/// %extracted_element = extract_element %val[%c0] : tensor<2xi32> +/// %extracted_element = tensor.extract %val[%c0] : tensor<2xi32> /// /// to /// -/// %extracted_element = extract_element %source[%c0] : tensor -struct ExtractElementFromTensorCast - : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; +/// %extracted_element = tensor.extract %source[%c0] : tensor +struct ExtractFromTensorCast : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(ExtractElementOp extract, + LogicalResult matchAndRewrite(tensor::ExtractOp extract, PatternRewriter &rewriter) const final { - auto tensorCast = extract.aggregate().getDefiningOp(); + auto tensorCast = extract.tensor().getDefiningOp(); if (!tensorCast) return failure(); - rewriter.replaceOpWithNewOp(extract, tensorCast.source(), - extract.getIndices()); + rewriter.replaceOpWithNewOp(extract, tensorCast.source(), + extract.indices()); return success(); } }; @@ -1919,51 +1920,9 @@ void DynamicTensorFromElementsOp::getCanonicalizationPatterns( OwningRewritePatternList &results, MLIRContext *context) { - results.insert( - context); -} - -//===----------------------------------------------------------------------===// -// ExtractElementOp -//===----------------------------------------------------------------------===// - -static LogicalResult verify(ExtractElementOp op) { - // Verify the # indices match if we have a ranked type. - auto aggregateType = op.getAggregate().getType().cast(); - if (aggregateType.hasRank() && - aggregateType.getRank() != op.getNumOperands() - 1) - return op.emitOpError("incorrect number of indices for extract_element"); - - return success(); -} - -OpFoldResult ExtractElementOp::fold(ArrayRef operands) { - assert(!operands.empty() && "extract_element takes at least one operand"); - - // The aggregate operand must be a known constant. - Attribute aggregate = operands.front(); - if (!aggregate) - return {}; - - // If this is a splat elements attribute, simply return the value. All of the - // elements of a splat attribute are the same. - if (auto splatAggregate = aggregate.dyn_cast()) - return splatAggregate.getSplatValue(); - - // Otherwise, collect the constant indices into the aggregate. - SmallVector indices; - for (Attribute indice : llvm::drop_begin(operands, 1)) { - if (!indice || !indice.isa()) - return {}; - indices.push_back(indice.cast().getInt()); - } - - // If this is an elements attribute, query the value at the given indices. - auto elementsAttr = aggregate.dyn_cast(); - if (elementsAttr && elementsAttr.isValidIndex(indices)) - return elementsAttr.getValue(indices); - return {}; + // TODO: Move extract patterns to tensor::ExtractOp. + results.insert(context); } //===----------------------------------------------------------------------===// @@ -1989,20 +1948,20 @@ // Canonicalizes the pattern of the form // // %tensor = "tensor_from_elements(%element) : (i32) -> tensor<1xi32> -// %extracted_element = extract_element %tensor[%c0] : tensor<1xi32> +// %extracted_element = tensor.extract %tensor[%c0] : tensor<1xi32> // // to just %element. struct ExtractElementFromTensorFromElements - : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(ExtractElementOp extract, + LogicalResult matchAndRewrite(tensor::ExtractOp extract, PatternRewriter &rewriter) const final { if (extract.indices().size() != 1) return failure(); auto tensorFromElements = dyn_cast_or_null( - extract.aggregate().getDefiningOp()); + extract.tensor().getDefiningOp()); if (tensorFromElements == nullptr) return failure(); @@ -2216,7 +2175,7 @@ } namespace { -/// Fold a load on a tensor_to_memref operation into an extract_element on the +/// Fold a load on a tensor_to_memref operation into an tensor.extract on the /// corresponding tensor. struct LoadOfTensorToMemref : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; @@ -2227,8 +2186,8 @@ if (!tensorToMemref) return failure(); - rewriter.replaceOpWithNewOp(load, tensorToMemref.tensor(), - load.indices()); + rewriter.replaceOpWithNewOp( + load, tensorToMemref.tensor(), load.indices()); return success(); } }; diff --git a/mlir/lib/Dialect/StandardOps/Transforms/Bufferize.cpp b/mlir/lib/Dialect/StandardOps/Transforms/Bufferize.cpp --- a/mlir/lib/Dialect/StandardOps/Transforms/Bufferize.cpp +++ b/mlir/lib/Dialect/StandardOps/Transforms/Bufferize.cpp @@ -15,6 +15,7 @@ #include "mlir/Dialect/SCF/SCF.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" #include "mlir/Dialect/StandardOps/Transforms/Passes.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/BlockAndValueMapping.h" #include "mlir/Transforms/DialectConversion.h" @@ -88,21 +89,6 @@ }; } // namespace -namespace { -class BufferizeExtractElementOp : public OpConversionPattern { -public: - using OpConversionPattern::OpConversionPattern; - LogicalResult - matchAndRewrite(ExtractElementOp op, ArrayRef operands, - ConversionPatternRewriter &rewriter) const override { - ExtractElementOp::Adaptor adaptor(operands); - rewriter.replaceOpWithNewOp(op, adaptor.aggregate(), - adaptor.indices()); - return success(); - } -}; -} // namespace - namespace { class BufferizeSelectOp : public OpConversionPattern { public: @@ -165,7 +151,6 @@ // clang-format off BufferizeDimOp, BufferizeDynamicTensorFromElementsOp, - BufferizeExtractElementOp, BufferizeSelectOp, BufferizeTensorCastOp, BufferizeTensorFromElementsOp @@ -183,10 +168,11 @@ target.addLegalDialect(); target.addLegalDialect(); + target.addLegalDialect(); populateStdBufferizePatterns(context, typeConverter, patterns); - target.addIllegalOp(); + target.addIllegalOp(); // We only bufferize the case of tensor selected type and scalar condition, // as that boils down to a select over memref descriptors (don't need to // touch the data). diff --git a/mlir/lib/Dialect/Tensor/CMakeLists.txt b/mlir/lib/Dialect/Tensor/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/Tensor/CMakeLists.txt @@ -0,0 +1,2 @@ +add_subdirectory(IR) +add_subdirectory(Transforms) diff --git a/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt b/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt @@ -0,0 +1,17 @@ +add_mlir_dialect_library(MLIRTensor + TensorDialect.cpp + TensorOps.cpp + + ADDITIONAL_HEADER_DIRS + ${PROJECT_SOURCE_DIR}/include/mlir/Dialect/Tensor + + DEPENDS + MLIRTensorOpsIncGen + + LINK_COMPONENTS + Core + + LINK_LIBS PUBLIC + MLIRIR + MLIRSupport + ) diff --git a/mlir/lib/Dialect/Tensor/IR/TensorDialect.cpp b/mlir/lib/Dialect/Tensor/IR/TensorDialect.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/Tensor/IR/TensorDialect.cpp @@ -0,0 +1,39 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Transforms/InliningUtils.h" + +using namespace mlir; +using namespace mlir::tensor; + +//===----------------------------------------------------------------------===// +// TensorDialect Dialect Interfaces +//===----------------------------------------------------------------------===// + +namespace { +struct TensorInlinerInterface : public DialectInlinerInterface { + using DialectInlinerInterface::DialectInlinerInterface; + bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned, + BlockAndValueMapping &valueMapping) const final { + return true; + } + bool isLegalToInline(Operation *, Region *, bool wouldBeCloned, + BlockAndValueMapping &) const final { + return true; + } +}; +} // end anonymous namespace + +void TensorDialect::initialize() { + addOperations< +#define GET_OP_LIST +#include "mlir/Dialect/Tensor/IR/TensorOps.cpp.inc" + >(); + addInterfaces(); +} diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp @@ -0,0 +1,60 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/TypeUtilities.h" +#include "llvm/ADT/STLExtras.h" + +using namespace mlir; +using namespace mlir::tensor; + +//===----------------------------------------------------------------------===// +// ExtractOp +//===----------------------------------------------------------------------===// + +static LogicalResult verify(ExtractOp op) { + // Verify the # indices match if we have a ranked type. + if (auto tensorType = op.tensor().getType().dyn_cast()) + if (tensorType.getRank() != static_cast(op.indices().size())) + return op.emitOpError("incorrect number of indices for extract_element"); + + return success(); +} + +OpFoldResult ExtractOp::fold(ArrayRef operands) { + // The tensor operand must be a known constant. + Attribute tensor = operands.front(); + if (!tensor) + return {}; + // If this is a splat elements attribute, simply return the value. All of the + // elements of a splat attribute are the same. + if (auto splatTensor = tensor.dyn_cast()) + return splatTensor.getSplatValue(); + + // Otherwise, collect the constant indices into the tensor. + SmallVector indices; + for (Attribute indice : llvm::drop_begin(operands, 1)) { + if (!indice || !indice.isa()) + return {}; + indices.push_back(indice.cast().getInt()); + } + + // If this is an elements attribute, query the value at the given indices. + auto elementsAttr = tensor.dyn_cast(); + if (elementsAttr && elementsAttr.isValidIndex(indices)) + return elementsAttr.getValue(indices); + return {}; +} + +//===----------------------------------------------------------------------===// +// TableGen'd op method definitions +//===----------------------------------------------------------------------===// + +#define GET_OP_CLASSES +#include "mlir/Dialect/Tensor/IR/TensorOps.cpp.inc" diff --git a/mlir/lib/Dialect/Tensor/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Tensor/Transforms/Bufferize.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/Tensor/Transforms/Bufferize.cpp @@ -0,0 +1,64 @@ +//===- Bufferize.cpp - Bufferization for `tensor` dialect ops -------------===// +// +// 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 bufferization of `tensor` dialect ops +// +//===----------------------------------------------------------------------===// + +#include "mlir/Transforms/Bufferize.h" +#include "PassDetail.h" +#include "mlir/Dialect/StandardOps/IR/Ops.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Dialect/Tensor/Transforms/Passes.h" +#include "mlir/Transforms/DialectConversion.h" + +using namespace mlir; + +namespace { +class BufferizeExtractOp : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(tensor::ExtractOp op, ArrayRef operands, + ConversionPatternRewriter &rewriter) const override { + tensor::ExtractOp::Adaptor adaptor(operands); + rewriter.replaceOpWithNewOp(op, adaptor.tensor(), + adaptor.indices()); + return success(); + } +}; +} // namespace + +void mlir::populateTensorBufferizePatterns( + MLIRContext *context, BufferizeTypeConverter &typeConverter, + OwningRewritePatternList &patterns) { + patterns.insert(typeConverter, context); +} + +namespace { +struct TensorBufferizePass : public TensorBufferizeBase { + void runOnFunction() override { + auto *context = &getContext(); + BufferizeTypeConverter typeConverter; + OwningRewritePatternList patterns; + ConversionTarget target(*context); + + populateTensorBufferizePatterns(context, typeConverter, patterns); + target.addIllegalOp(); + target.addLegalDialect(); + + if (failed( + applyPartialConversion(getFunction(), target, std::move(patterns)))) + signalPassFailure(); + } +}; +} // namespace + +std::unique_ptr mlir::createTensorBufferizePass() { + return std::make_unique(); +} diff --git a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt @@ -0,0 +1,15 @@ +add_mlir_dialect_library(MLIRTensorTransforms + Bufferize.cpp + + ADDITIONAL_HEADER_DIRS + ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tensor/Transforms + + DEPENDS + MLIRTensorTransformsIncGen + + LINK_LIBS PUBLIC + MLIRIR + MLIRPass + MLIRTensor + MLIRTransforms + ) diff --git a/mlir/lib/Dialect/Tensor/Transforms/PassDetail.h b/mlir/lib/Dialect/Tensor/Transforms/PassDetail.h new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/Tensor/Transforms/PassDetail.h @@ -0,0 +1,21 @@ +//===- PassDetail.h - GPU Pass class details --------------------*- C++ -*-===// +// +// 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 DIALECT_TENSOR_TRANSFORMS_PASSDETAIL_H_ +#define DIALECT_TENSOR_TRANSFORMS_PASSDETAIL_H_ + +#include "mlir/Pass/Pass.h" + +namespace mlir { + +#define GEN_PASS_CLASSES +#include "mlir/Dialect/Tensor/Transforms/Passes.h.inc" + +} // end namespace mlir + +#endif // DIALECT_TENSOR_TRANSFORMS_PASSDETAIL_H_ diff --git a/mlir/lib/Transforms/Utils/FoldUtils.cpp b/mlir/lib/Transforms/Utils/FoldUtils.cpp --- a/mlir/lib/Transforms/Utils/FoldUtils.cpp +++ b/mlir/lib/Transforms/Utils/FoldUtils.cpp @@ -13,6 +13,7 @@ #include "mlir/Transforms/FoldUtils.h" +#include "mlir/Dialect/StandardOps/IR/Ops.h" #include "mlir/IR/Builders.h" #include "mlir/IR/Matchers.h" #include "mlir/IR/Operation.h" @@ -59,6 +60,23 @@ assert(matchPattern(constOp, m_Constant())); return constOp; } + + // TODO: To faciliate splitting the std dialect (PR48490), have a special case + // for falling back to std.constant. Eventually, we will have separate ops + // tensor.constant, int.constant, float.constant, etc. that live in their + // respective dialects, which will allow each dialect to implement the + // materializeConstant hook above. + // + // The special case is needed because in the interim state while we are + // splitting out those dialects from std, the std dialect depends on the + // tensor dialect, which makes it impossible for the tensor dialect to use + // std.constant (it would be a cyclic dependency) as part of its + // materializeConstant hook. + // + // If the dialect is unable to materialize a constant, check to see if the + // standard constant can be used. + if (ConstantOp::isBuildableWith(value, type)) + return builder.create(loc, type, value); return nullptr; } diff --git a/mlir/test/Conversion/ShapeToStandard/convert-shape-constraints.mlir b/mlir/test/Conversion/ShapeToStandard/convert-shape-constraints.mlir --- a/mlir/test/Conversion/ShapeToStandard/convert-shape-constraints.mlir +++ b/mlir/test/Conversion/ShapeToStandard/convert-shape-constraints.mlir @@ -16,9 +16,9 @@ // CHECK: %[[GREATER_RANK_OPERAND:.*]] = select %[[LHS_RANK_ULE]], %[[RHS]], %[[LHS]] : tensor // CHECK: %[[RANK_DIFF:.*]] = subi %[[GREATER_RANK]], %[[LESSER_RANK]] : index // CHECK: scf.for %[[IV:.*]] = %[[RANK_DIFF]] to %[[GREATER_RANK]] step %[[C1]] { -// CHECK: %[[GREATER_RANK_OPERAND_EXTENT:.*]] = extract_element %[[GREATER_RANK_OPERAND]][%[[IV]]] : tensor +// CHECK: %[[GREATER_RANK_OPERAND_EXTENT:.*]] = tensor.extract %[[GREATER_RANK_OPERAND]][%[[IV]]] : tensor // CHECK: %[[IVSHIFTED:.*]] = subi %[[IV]], %[[RANK_DIFF]] : index -// CHECK: %[[LESSER_RANK_OPERAND_EXTENT:.*]] = extract_element %[[LESSER_RANK_OPERAND]][%[[IVSHIFTED]]] : tensor +// CHECK: %[[LESSER_RANK_OPERAND_EXTENT:.*]] = tensor.extract %[[LESSER_RANK_OPERAND]][%[[IVSHIFTED]]] : tensor // CHECK: %[[GREATER_RANK_OPERAND_EXTENT_IS_ONE:.*]] = cmpi "eq", %[[GREATER_RANK_OPERAND_EXTENT]], %[[C1]] : index // CHECK: %[[LESSER_RANK_OPERAND_EXTENT_IS_ONE:.*]] = cmpi "eq", %[[LESSER_RANK_OPERAND_EXTENT]], %[[C1]] : index // CHECK: %[[EXTENTS_AGREE:.*]] = cmpi "eq", %[[GREATER_RANK_OPERAND_EXTENT]], %[[LESSER_RANK_OPERAND_EXTENT]] : index diff --git a/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir b/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir --- a/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir +++ b/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir @@ -74,12 +74,12 @@ // ----- -// Express `get_extent` as `std.extract_element`. +// Express `get_extent` as `std.tensor.extract`. // CHECK-LABEL: @get_extent_from_extent_tensor // CHECK-SAME: (%[[EXTENTS:.*]]: tensor, %[[IDX:.*]]: index) -> index func @get_extent_from_extent_tensor(%extents : tensor, %idx : index) -> index { - // CHECK: %[[RESULT:.*]] = extract_element %[[EXTENTS]][%[[IDX]]] : tensor + // CHECK: %[[RESULT:.*]] = tensor.extract %[[EXTENTS]][%[[IDX]]] : tensor // CHECK: return %[[RESULT]] : index %result = shape.get_extent %extents, %idx : tensor, index -> index return %result : index @@ -180,7 +180,7 @@ // CHECK-NEXT: %[[C1:.*]] = constant 1 : index // CHECK-NEXT: %[[RANK:.*]] = dim %[[SHAPE]], %[[C0]] : tensor // CHECK-NEXT: %[[RESULT:.*]] = scf.for %[[I:.*]] = %[[C0]] to %[[RANK]] step %[[C1]] iter_args(%[[ACC:.*]] = %[[INIT]]) -> (index) -// CHECK-NEXT: %[[EXTENT:.*]] = extract_element %[[SHAPE]][%[[I]]] +// CHECK-NEXT: %[[EXTENT:.*]] = tensor.extract %[[SHAPE]][%[[I]]] // CHECK-NEXT: %[[NEW_ACC:.*]] = muli %[[ACC]], %[[EXTENT]] : index // CHECK-NEXT: scf.yield %[[NEW_ACC]] : index // CHECK-NEXT: } @@ -277,8 +277,8 @@ // CHECK: %[[C1:.*]] = constant 1 : index // CHECK: %[[INIT:.*]] = constant true // CHECK: %[[SHAPE_EQ_INNER:.*]] = scf.for %[[I:.*]] = %[[C0]] to %[[RANK_A]] step %[[C1]] iter_args(%[[CONJ:.*]] = %[[INIT]]) -> (i1) { - // CHECK: %[[EXTENT_A:.*]] = extract_element %[[A]][%[[I]]] : tensor - // CHECK: %[[EXTENT_B:.*]] = extract_element %[[B]][%[[I]]] : tensor + // CHECK: %[[EXTENT_A:.*]] = tensor.extract %[[A]][%[[I]]] : tensor + // CHECK: %[[EXTENT_B:.*]] = tensor.extract %[[B]][%[[I]]] : tensor // CHECK: %[[EXTENT_EQ:.*]] = cmpi "eq", %[[EXTENT_A]], %[[EXTENT_B]] // CHECK: %[[CONJ_NEXT:.*]] = and %[[CONJ]], %[[EXTENT_EQ]] // CHECK: scf.yield %[[CONJ_NEXT]] : i1 @@ -324,12 +324,12 @@ // CHECK: %[[RESULT:.*]] = dynamic_tensor_from_elements %[[GREATER_RANK]] { // CHECK: ^bb0(%[[OUTPUT_DIMENSION:.*]]: index): // CHECK: %[[IS_UNCHALLENGED_DIMENSION:.*]] = cmpi "ult", %[[OUTPUT_DIMENSION]], %[[RANK_DIFF]] : index - // CHECK: %[[GREATER_RANK_OPERAND_EXTENT:.*]] = extract_element %[[GREATER_RANK_OPERAND]][%[[OUTPUT_DIMENSION]]] : tensor + // CHECK: %[[GREATER_RANK_OPERAND_EXTENT:.*]] = tensor.extract %[[GREATER_RANK_OPERAND]][%[[OUTPUT_DIMENSION]]] : tensor // CHECK: %[[OUTPUT_EXTENT:.*]] = scf.if %[[IS_UNCHALLENGED_DIMENSION]] -> (index) { // CHECK: scf.yield %[[GREATER_RANK_OPERAND_EXTENT]] : index // CHECK: } else { // CHECK: %[[LESSER_RANK_OPERAND_DIMENSION:.*]] = subi %[[OUTPUT_DIMENSION]], %[[RANK_DIFF]] : index - // CHECK: %[[LESSER_RANK_OPERAND_EXTENT:.*]] = extract_element %[[LESSER_RANK_OPERAND]][%[[LESSER_RANK_OPERAND_DIMENSION]]] : tensor + // CHECK: %[[LESSER_RANK_OPERAND_EXTENT:.*]] = tensor.extract %[[LESSER_RANK_OPERAND]][%[[LESSER_RANK_OPERAND_DIMENSION]]] : tensor // CHECK: %[[GREATER_RANK_OPERAND_EXTENT_IS_ONE:.*]] = cmpi "eq", %[[GREATER_RANK_OPERAND_EXTENT]], %[[C1]] : index // CHECK: %[[BROADCASTED_EXTENT:.*]] = select %[[GREATER_RANK_OPERAND_EXTENT_IS_ONE]], %[[LESSER_RANK_OPERAND_EXTENT]], %[[GREATER_RANK_OPERAND_EXTENT]] : index // CHECK: scf.yield %[[BROADCASTED_EXTENT]] : index @@ -364,12 +364,12 @@ // CHECK: %[[RESULT:.*]] = dynamic_tensor_from_elements %[[GREATER_RANK]] { // CHECK: ^bb0(%[[OUTPUT_DIMENSION:.*]]: index): // CHECK: %[[IS_UNCHALLENGED_DIMENSION:.*]] = cmpi "ult", %[[OUTPUT_DIMENSION]], %[[RANK_DIFF]] : index - // CHECK: %[[GREATER_RANK_OPERAND_EXTENT:.*]] = extract_element %[[GREATER_RANK_OPERAND]][%[[OUTPUT_DIMENSION]]] : tensor + // CHECK: %[[GREATER_RANK_OPERAND_EXTENT:.*]] = tensor.extract %[[GREATER_RANK_OPERAND]][%[[OUTPUT_DIMENSION]]] : tensor // CHECK: %[[OUTPUT_EXTENT:.*]] = scf.if %[[IS_UNCHALLENGED_DIMENSION]] -> (index) { // CHECK: scf.yield %[[GREATER_RANK_OPERAND_EXTENT]] : index // CHECK: } else { // CHECK: %[[LESSER_RANK_OPERAND_DIMENSION:.*]] = subi %[[OUTPUT_DIMENSION]], %[[RANK_DIFF]] : index - // CHECK: %[[LESSER_RANK_OPERAND_EXTENT:.*]] = extract_element %[[LESSER_RANK_OPERAND]][%[[LESSER_RANK_OPERAND_DIMENSION]]] : tensor + // CHECK: %[[LESSER_RANK_OPERAND_EXTENT:.*]] = tensor.extract %[[LESSER_RANK_OPERAND]][%[[LESSER_RANK_OPERAND_DIMENSION]]] : tensor // CHECK: %[[GREATER_RANK_OPERAND_EXTENT_IS_ONE:.*]] = cmpi "eq", %[[GREATER_RANK_OPERAND_EXTENT]], %[[C1]] : index // CHECK: %[[BROADCASTED_EXTENT:.*]] = select %[[GREATER_RANK_OPERAND_EXTENT_IS_ONE]], %[[LESSER_RANK_OPERAND_EXTENT]], %[[GREATER_RANK_OPERAND_EXTENT]] : index // CHECK: scf.yield %[[BROADCASTED_EXTENT]] : index @@ -407,10 +407,10 @@ // CHECK: %[[RANK_DIFF:.*]] = subi %[[LARGER_RANK]], %[[SMALLER_RANK]] : index // CHECK: %[[TRUE:.*]] = constant true // CHECK: %[[ALL_RESULT:.*]] = scf.for %[[I:.*]] = %[[RANK_DIFF]] to %[[LARGER_RANK]] step %[[C1]] iter_args(%[[ALL_SO_FAR:.*]] = %[[TRUE]]) -> (i1) { -// CHECK: %[[LARGER_EXTENT:.*]] = extract_element %[[LARGER_SHAPE]]{{\[}}%[[I]]] : tensor +// CHECK: %[[LARGER_EXTENT:.*]] = tensor.extract %[[LARGER_SHAPE]]{{\[}}%[[I]]] : tensor // CHECK: %[[LARGER_EXTENT_IS_ONE:.*]] = cmpi "eq", %[[LARGER_EXTENT]], %[[C1]] : index // CHECK: %[[SMALLER_EXTENT_INDEX:.*]] = subi %[[I]], %[[RANK_DIFF]] : index -// CHECK: %[[SMALLER_EXTENT:.*]] = extract_element %[[SMALLER_SHAPE]]{{\[}}%[[SMALLER_EXTENT_INDEX]]] : tensor +// CHECK: %[[SMALLER_EXTENT:.*]] = tensor.extract %[[SMALLER_SHAPE]]{{\[}}%[[SMALLER_EXTENT_INDEX]]] : tensor // CHECK: %[[SMALLER_EXTENT_IS_ONE:.*]] = cmpi "eq", %[[SMALLER_EXTENT]], %[[C1]] : index // CHECK: %[[EXTENTS_ARE_EQUAL:.*]] = cmpi "eq", %[[LARGER_EXTENT]], %[[SMALLER_EXTENT]] : index // CHECK: %[[EITHER_EXTENT_IS_ONE:.*]] = or %[[LARGER_EXTENT_IS_ONE]], %[[SMALLER_EXTENT_IS_ONE]] : i1 @@ -445,10 +445,10 @@ // CHECK: %[[RANK_DIFF:.*]] = subi %[[LARGER_RANK]], %[[SMALLER_RANK]] : index // CHECK: %[[TRUE:.*]] = constant true // CHECK: %[[ALL_RESULT:.*]] = scf.for %[[VAL_16:.*]] = %[[RANK_DIFF]] to %[[LARGER_RANK]] step %[[C1]] iter_args(%[[ALL_SO_FAR:.*]] = %[[TRUE]]) -> (i1) { -// CHECK: %[[LARGER_EXTENT:.*]] = extract_element %[[LARGER_SHAPE]]{{\[}}%[[VAL_16]]] : tensor +// CHECK: %[[LARGER_EXTENT:.*]] = tensor.extract %[[LARGER_SHAPE]]{{\[}}%[[VAL_16]]] : tensor // CHECK: %[[LARGER_EXTENT_IS_ONE:.*]] = cmpi "eq", %[[LARGER_EXTENT]], %[[C1]] : index // CHECK: %[[LHS_EXTENT_INDEX:.*]] = subi %[[VAL_16]], %[[RANK_DIFF]] : index -// CHECK: %[[SMALLER_EXTENT:.*]] = extract_element %[[SMALLER_SHAPE]]{{\[}}%[[LHS_EXTENT_INDEX]]] : tensor +// CHECK: %[[SMALLER_EXTENT:.*]] = tensor.extract %[[SMALLER_SHAPE]]{{\[}}%[[LHS_EXTENT_INDEX]]] : tensor // CHECK: %[[SMALLER_EXTENT_IS_ONE:.*]] = cmpi "eq", %[[SMALLER_EXTENT]], %[[C1]] : index // CHECK: %[[EXTENTS_ARE_EQUAL:.*]] = cmpi "eq", %[[LARGER_EXTENT]], %[[SMALLER_EXTENT]] : index // CHECK: %[[EITHER_EXTENT_IS_ONE:.*]] = or %[[LARGER_EXTENT_IS_ONE]], %[[SMALLER_EXTENT_IS_ONE]] : i1 diff --git a/mlir/test/Dialect/Linalg/fusion-tensor.mlir b/mlir/test/Dialect/Linalg/fusion-tensor.mlir --- a/mlir/test/Dialect/Linalg/fusion-tensor.mlir +++ b/mlir/test/Dialect/Linalg/fusion-tensor.mlir @@ -395,7 +395,7 @@ ins(%arg1 : tensor) { ^bb0(%arg2: i32): // no predecessors %3 = index_cast %arg2 : i32 to index - %4 = extract_element %arg0[%3, %c0, %c0] : tensor<5x1x1xf32> + %4 = tensor.extract %arg0[%3, %c0, %c0] : tensor<5x1x1xf32> linalg.yield %4 : f32 } -> tensor %1 = linalg.generic @@ -418,6 +418,6 @@ // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]] // CHECK-SAME: iterator_types = ["parallel"] // CHECK-SAME: ins(%[[ARG1]] : tensor) -// CHECK: extract_element %[[ARG0]] +// CHECK: tensor.extract %[[ARG0]] // CHECK: linalg.yield // CHECK return %[[T0]] diff --git a/mlir/test/Dialect/Standard/bufferize.mlir b/mlir/test/Dialect/Standard/bufferize.mlir --- a/mlir/test/Dialect/Standard/bufferize.mlir +++ b/mlir/test/Dialect/Standard/bufferize.mlir @@ -61,18 +61,6 @@ return %result : tensor<16x?xindex> } -// CHECK-LABEL: func @extract_element( -// CHECK-SAME: %[[TENSOR:.*]]: tensor, -// CHECK-SAME: %[[IDX:.*]]: index) -> f32 { -// CHECK: %[[MEMREF:.*]] = tensor_to_memref %[[TENSOR]] : memref -// CHECK: %[[RET:.*]] = load %[[MEMREF]][%[[IDX]]] : memref -// CHECK: return %[[RET]] : f32 -// CHECK: } -func @extract_element(%arg0: tensor, %arg1: index) -> f32 { - %0 = extract_element %arg0[%arg1] : tensor - return %0 : f32 -} - // CHECK-LABEL: func @select( // CHECK-SAME: %[[PRED:.*]]: i1, // CHECK-SAME: %[[TRUE_VAL:.*]]: tensor, @@ -138,14 +126,14 @@ // The dynamic_tensor_from_elements op clones each op in its body. // Make sure that regions nested within such ops are recursively converted. // CHECK-LABEL: func @recursively_convert_cloned_regions -func @recursively_convert_cloned_regions(%arg0: tensor, %arg1: index, %arg2: i1) -> tensor { +func @recursively_convert_cloned_regions(%arg0: tensor<*xf32>, %arg1: index, %arg2: i1) -> tensor { %tensor = dynamic_tensor_from_elements %arg1 { ^bb0(%iv: index): %48 = scf.if %arg2 -> (index) { scf.yield %iv : index } else { - // CHECK-NOT: extract_element - %50 = extract_element %arg0[%iv] : tensor + // CHECK-NOT: dim{{.*}}tensor + %50 = dim %arg0, %iv : tensor<*xf32> scf.yield %50 : index } yield %48 : index diff --git a/mlir/test/Dialect/Standard/canonicalize.mlir b/mlir/test/Dialect/Standard/canonicalize.mlir --- a/mlir/test/Dialect/Standard/canonicalize.mlir +++ b/mlir/test/Dialect/Standard/canonicalize.mlir @@ -46,11 +46,11 @@ } // Test case: Folding of load(tensor_to_memref(%v, %idxs)) -// -> extract_element(%v, %idx) +// -> tensor.extract(%v, %idx) // CHECK-LABEL: func @load_from_tensor_to_memref( // CHECK-SAME: %[[IDX0:[0-9a-z]+]]: index, %[[IDX1:[0-9a-z]+]]: index // CHECK-SAME: %[[TENSOR:[0-9a-z]+]]: tensor -// CHECK: %[[RES:.*]] = extract_element %[[TENSOR]][%[[IDX0]], %[[IDX1]]] +// CHECK: %[[RES:.*]] = tensor.extract %[[TENSOR]][%[[IDX0]], %[[IDX1]]] // CHECK-NOT: load // CHECK: return %[[RES]] : f32 func @load_from_tensor_to_memref(%arg0: index, %arg1: index, %arg2: tensor) -> f32 { diff --git a/mlir/test/Dialect/Tensor/bufferize.mlir b/mlir/test/Dialect/Tensor/bufferize.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/Tensor/bufferize.mlir @@ -0,0 +1,13 @@ +// RUN: mlir-opt %s -tensor-bufferize | FileCheck %s + +// CHECK-LABEL: func @extract( +// CHECK-SAME: %[[TENSOR:.*]]: tensor, +// CHECK-SAME: %[[IDX:.*]]: index) -> f32 { +// CHECK: %[[MEMREF:.*]] = tensor_to_memref %[[TENSOR]] : memref +// CHECK: %[[RET:.*]] = load %[[MEMREF]][%[[IDX]]] : memref +// CHECK: return %[[RET]] : f32 +// CHECK: } +func @extract(%arg0: tensor, %arg1: index) -> f32 { + %0 = tensor.extract %arg0[%arg1] : tensor + return %0 : f32 +} diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/Tensor/canonicalize.mlir @@ -0,0 +1,33 @@ +// RUN: mlir-opt %s -canonicalize | FileCheck %s + +// ----- + +// CHECK-LABEL: func @fold_extract +func @fold_extract(%arg0 : index) -> (f32, f16, f16, i32) { + %const_0 = constant 0 : index + %const_1 = constant 1 : index + %const_3 = constant 3 : index + + // Fold an extract into a splat. + // CHECK-NEXT: [[C4:%.+]] = constant 4.{{0*}}e+00 : f32 + %0 = constant dense<4.0> : tensor<4xf32> + %ext_1 = tensor.extract %0[%arg0] : tensor<4xf32> + + // Fold an extract into a sparse with a sparse index. + // CHECK-NEXT: [[CM2:%.+]] = constant -2.{{0*}}e+00 : f16 + %1 = constant sparse<[[0, 0, 0], [1, 1, 1]], [-5.0, -2.0]> : tensor<4x4x4xf16> + %ext_2 = tensor.extract %1[%const_1, %const_1, %const_1] : tensor<4x4x4xf16> + + // Fold an extract into a sparse with a non sparse index. + // CHECK-NEXT: [[C0:%.+]] = constant 0.{{0*}}e+00 : f16 + %2 = constant sparse<[[1, 1, 1]], [-2.0]> : tensor<1x1x1xf16> + %ext_3 = tensor.extract %2[%const_0, %const_0, %const_0] : tensor<1x1x1xf16> + + // Fold an extract into a dense tensor. + // CHECK-NEXT: [[C64:%.+]] = constant 64 : i32 + %3 = constant dense<[[[1, -2, 1, 36]], [[0, 2, -1, 64]]]> : tensor<2x1x4xi32> + %ext_4 = tensor.extract %3[%const_1, %const_0, %const_3] : tensor<2x1x4xi32> + + // CHECK-NEXT: return [[C4]], [[CM2]], [[C0]], [[C64]] + return %ext_1, %ext_2, %ext_3, %ext_4 : f32, f16, f16, i32 +} diff --git a/mlir/test/Dialect/Tensor/invalid.mlir b/mlir/test/Dialect/Tensor/invalid.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/Tensor/invalid.mlir @@ -0,0 +1,9 @@ +// RUN: mlir-opt <%s -verify-diagnostics + +// ----- + +func @extract_too_many_indices(%arg0: tensor) { + // expected-error@+1 {{incorrect number of indices for extract_element}} + %0 = tensor.extract %arg0[] : tensor + return +} diff --git a/mlir/test/Dialect/Tensor/ops.mlir b/mlir/test/Dialect/Tensor/ops.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/Tensor/ops.mlir @@ -0,0 +1,10 @@ +// RUN: mlir-opt <%s | mlir-opt | FileCheck %s + +// CHECK-LABEL: func @extract( +// CHECK-SAME: %[[TENSOR:.*]]: tensor, +// CHECK-SAME: %[[INDEX:.*]]: index) { +func @extract(%arg0: tensor, %arg1: index) { + // CHECK: tensor.extract %[[TENSOR]][%[[INDEX]], %[[INDEX]], %[[INDEX]]] : tensor + %0 = tensor.extract %arg0[%arg1, %arg1, %arg1] : tensor + return +} diff --git a/mlir/test/IR/core-ops.mlir b/mlir/test/IR/core-ops.mlir --- a/mlir/test/IR/core-ops.mlir +++ b/mlir/test/IR/core-ops.mlir @@ -672,19 +672,6 @@ return } -// CHECK-LABEL: func @extract_element(%arg0: tensor<*xi32>, %arg1: tensor<4x4xf32>) -> i32 { -func @extract_element(%arg0: tensor<*xi32>, %arg1 : tensor<4x4xf32>) -> i32 { - %c0 = "std.constant"() {value = 0: index} : () -> index - - // CHECK: %0 = extract_element %arg0[%c0, %c0, %c0, %c0] : tensor<*xi32> - %0 = extract_element %arg0[%c0, %c0, %c0, %c0] : tensor<*xi32> - - // CHECK: %1 = extract_element %arg1[%c0, %c0] : tensor<4x4xf32> - %1 = extract_element %arg1[%c0, %c0] : tensor<4x4xf32> - - return %0 : i32 -} - // CHECK-LABEL: func @tensor_from_elements() { func @tensor_from_elements() { %c0 = "std.constant"() {value = 0: index} : () -> index @@ -972,4 +959,3 @@ // CHECK-LABEL: func private @legacy_visibility_syntax func @legacy_visibility_syntax() attributes { sym_visibility = "private" } - diff --git a/mlir/test/IR/invalid-ops.mlir b/mlir/test/IR/invalid-ops.mlir --- a/mlir/test/IR/invalid-ops.mlir +++ b/mlir/test/IR/invalid-ops.mlir @@ -541,61 +541,6 @@ // ----- -func @extract_element_no_operands() { - // expected-error@+1 {{op expected 1 or more operands}} - %0 = "std.extract_element"() : () -> f32 - return -} - -// ----- - -func @extract_element_no_indices(%v : vector<3xf32>) { - // expected-error@+1 {{incorrect number of indices for extract_element}} - %0 = "std.extract_element"(%v) : (vector<3xf32>) -> f32 - return -} - -// ----- - -func @extract_element_invalid_index_type(%v : vector<3xf32>, %i : i32) { - // expected-error@+1 {{operand #1 must be index}} - %0 = "std.extract_element"(%v, %i) : (vector<3xf32>, i32) -> f32 - return -} - -// ----- - -func @extract_element_element_result_type_mismatch(%v : vector<3xf32>, %i : index) { - // expected-error@+1 {{result type matches element type of aggregate}} - %0 = "std.extract_element"(%v, %i) : (vector<3xf32>, index) -> f64 - return -} - -// ----- - -func @extract_element_vector_too_many_indices(%v : vector<3xf32>, %i : index) { - // expected-error@+1 {{incorrect number of indices for extract_element}} - %0 = "std.extract_element"(%v, %i, %i) : (vector<3xf32>, index, index) -> f32 - return -} - -// ----- - -func @extract_element_tensor_too_many_indices(%t : tensor<2x3xf32>, %i : index) { - // expected-error@+1 {{incorrect number of indices for extract_element}} - %0 = "std.extract_element"(%t, %i, %i, %i) : (tensor<2x3xf32>, index, index, index) -> f32 - return -} - -// ----- - -func @extract_element_tensor_too_few_indices(%t : tensor<2x3xf32>, %i : index) { - // expected-error@+1 {{incorrect number of indices for extract_element}} - %0 = "std.extract_element"(%t, %i) : (tensor<2x3xf32>, index) -> f32 return -} - -// ----- - func @tensor_from_elements_wrong_result_type() { // expected-error@+2 {{'result' must be 1D tensor of any type values, but got 'tensor<*xi32>'}} %c0 = constant 0 : i32 diff --git a/mlir/test/Transforms/canonicalize.mlir b/mlir/test/Transforms/canonicalize.mlir --- a/mlir/test/Transforms/canonicalize.mlir +++ b/mlir/test/Transforms/canonicalize.mlir @@ -1040,21 +1040,21 @@ // ----- -// CHECK-LABEL: func @extract_element_from_tensor_from_elements -func @extract_element_from_tensor_from_elements(%element : index) -> index { +// CHECK-LABEL: func @extract_from_tensor_from_elements +func @extract_from_tensor_from_elements(%element : index) -> index { // CHECK-SAME: ([[ARG:%.*]]: index) %c0 = constant 0 : index %tensor = tensor_from_elements %element : tensor<1xindex> - %extracted_element = extract_element %tensor[%c0] : tensor<1xindex> + %extracted_element = tensor.extract %tensor[%c0] : tensor<1xindex> // CHECK: [[ARG]] : index return %extracted_element : index } // ----- -// CHECK-LABEL: func @extract_element_from_dynamic_tensor_from_elements +// CHECK-LABEL: func @extract_from_dynamic_tensor_from_elements // CHECK-SAME: %[[IDX:.*]]: index, %[[TENSOR:.*]]: tensor<*xf32> -func @extract_element_from_dynamic_tensor_from_elements(%idx: index, %tensor: tensor<*xf32>) -> index { +func @extract_from_dynamic_tensor_from_elements(%idx: index, %tensor: tensor<*xf32>) -> index { %size = rank %tensor : tensor<*xf32> // CHECK-NEXT: %[[RES:.*]] = dim %[[TENSOR]], %[[IDX]] %0 = dynamic_tensor_from_elements %size { @@ -1062,16 +1062,16 @@ %1 = dim %tensor, %arg0 : tensor<*xf32> yield %1 : index } : tensor - %1 = extract_element %0[%idx] : tensor + %1 = tensor.extract %0[%idx] : tensor // CHECK-NEXT: return %[[RES]] return %1 : index } // ----- -// CHECK-LABEL: func @extract_element_from_dynamic_tensor_from_elements_2d +// CHECK-LABEL: func @extract_from_dynamic_tensor_from_elements_2d // CHECK-SAME: %[[IDX0:.*]]: index, %[[IDX1:.*]]: index, %[[TENSOR:.*]]: tensor<*xf32> -func @extract_element_from_dynamic_tensor_from_elements_2d(%idx0: index, %idx1: index, %tensor: tensor<*xf32>) -> index { +func @extract_from_dynamic_tensor_from_elements_2d(%idx0: index, %idx1: index, %tensor: tensor<*xf32>) -> index { %size = rank %tensor : tensor<*xf32> // CHECK-NEXT: %[[DIM0:.*]] = dim %[[TENSOR]], %[[IDX0]] // CHECK-NEXT: %[[DIM1:.*]] = dim %[[TENSOR]], %[[IDX1]] @@ -1083,16 +1083,16 @@ %3 = addi %1, %2 : index yield %3 : index } : tensor - %4 = extract_element %0[%idx0, %idx1] : tensor + %4 = tensor.extract %0[%idx0, %idx1] : tensor // CHECK-NEXT: return %[[RES]] return %4 : index } // ----- -// CHECK-LABEL: func @extract_element_from_dynamic_tensor_from_elements_sideeffects +// CHECK-LABEL: func @extract_from_dynamic_tensor_from_elements_sideeffects // CHECK-SAME: %[[IDX:.*]]: index -func @extract_element_from_dynamic_tensor_from_elements_sideeffects(%idx: index, %tensor: tensor<*xf32>) -> index { +func @extract_from_dynamic_tensor_from_elements_sideeffects(%idx: index, %tensor: tensor<*xf32>) -> index { %size = rank %tensor : tensor<*xf32> %mem = alloc(%size) : memref // CHECK: %[[DTENSOR:.*]] = dynamic_tensor_from_elements @@ -1102,8 +1102,8 @@ store %1, %mem[%arg0] : memref yield %1 : index } : tensor - // CHECK: %[[RES:.*]] = extract_element %[[DTENSOR]][%[[IDX]]] - %1 = extract_element %0[%idx] : tensor + // CHECK: %[[RES:.*]] = tensor.extract %[[DTENSOR]][%[[IDX]]] + %1 = tensor.extract %0[%idx] : tensor // CHECK-NEXT: return %[[RES]] return %1 : index } @@ -1205,14 +1205,14 @@ // ----- -// CHECK-LABEL: func @extract_element_from_tensor_cast +// CHECK-LABEL: func @extract_from_tensor_cast // CHECK-SAME: %[[TENSOR:.*]]: tensor<*xf32> -func @extract_element_from_tensor_cast(%tensor: tensor<*xf32>) -> f32 { +func @extract_from_tensor_cast(%tensor: tensor<*xf32>) -> f32 { // CHECK-NEXT: %[[C0:.*]] = constant 0 : index %c0 = constant 0 : index // CHECK-NOT: tensor_cast %casted = tensor_cast %tensor : tensor<*xf32> to tensor - // CHECK-NEXT: extract_element %[[TENSOR]][%[[C0]]] - %result = extract_element %casted[%c0] : tensor + // CHECK-NEXT: tensor.extract %[[TENSOR]][%[[C0]]] + %result = tensor.extract %casted[%c0] : tensor return %result : f32 } diff --git a/mlir/test/Transforms/constant-fold.mlir b/mlir/test/Transforms/constant-fold.mlir --- a/mlir/test/Transforms/constant-fold.mlir +++ b/mlir/test/Transforms/constant-fold.mlir @@ -716,38 +716,6 @@ // ----- -// CHECK-LABEL: func @fold_extract_element -func @fold_extract_element(%arg0 : index) -> (f32, f16, f16, i32) { - %const_0 = constant 0 : index - %const_1 = constant 1 : index - %const_3 = constant 3 : index - - // Fold an extract into a splat. - // CHECK-NEXT: [[C4:%.+]] = constant 4.{{0*}}e+00 : f32 - %0 = constant dense<4.0> : tensor<4xf32> - %ext_1 = extract_element %0[%arg0] : tensor<4xf32> - - // Fold an extract into a sparse with a sparse index. - // CHECK-NEXT: [[CM2:%.+]] = constant -2.{{0*}}e+00 : f16 - %1 = constant sparse<[[0, 0, 0], [1, 1, 1]], [-5.0, -2.0]> : vector<4x4x4xf16> - %ext_2 = extract_element %1[%const_1, %const_1, %const_1] : vector<4x4x4xf16> - - // Fold an extract into a sparse with a non sparse index. - // CHECK-NEXT: [[C0:%.+]] = constant 0.{{0*}}e+00 : f16 - %2 = constant sparse<[[1, 1, 1]], [-2.0]> : vector<1x1x1xf16> - %ext_3 = extract_element %2[%const_0, %const_0, %const_0] : vector<1x1x1xf16> - - // Fold an extract into a dense tensor. - // CHECK-NEXT: [[C64:%.+]] = constant 64 : i32 - %3 = constant dense<[[[1, -2, 1, 36]], [[0, 2, -1, 64]]]> : tensor<2x1x4xi32> - %ext_4 = extract_element %3[%const_1, %const_0, %const_3] : tensor<2x1x4xi32> - - // CHECK-NEXT: return [[C4]], [[CM2]], [[C0]], [[C64]] - return %ext_1, %ext_2, %ext_3, %ext_4 : f32, f16, f16, i32 -} - -// ----- - // CHECK-LABEL: func @fold_rank func @fold_rank() -> (index) { %const_0 = constant dense<[[[1, -2, 1, 36]], [[0, 2, -1, 64]]]> : tensor<2x1x4xi32> diff --git a/mlir/utils/vim/syntax/mlir.vim b/mlir/utils/vim/syntax/mlir.vim --- a/mlir/utils/vim/syntax/mlir.vim +++ b/mlir/utils/vim/syntax/mlir.vim @@ -38,7 +38,7 @@ " TODO: this list is not exhaustive. syn keyword mlirOps alloc alloca addf addi and call call_indirect cmpf cmpi syn keyword mlirOps constant dealloc divf dma_start dma_wait dim exp -syn keyword mlirOps extract_element getTensor index_cast load log memref_cast +syn keyword mlirOps getTensor index_cast load log memref_cast syn keyword mlirOps memref_shape_cast mulf muli negf powf prefetch rsqrt sitofp syn keyword mlirOps splat store select sqrt subf subi subview tanh tensor_cast syn keyword mlirOps view