diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td --- a/mlir/include/mlir/Conversion/Passes.td +++ b/mlir/include/mlir/Conversion/Passes.td @@ -367,7 +367,8 @@ let dependentDialects = [ "memref::MemRefDialect", "StandardOpsDialect", - "scf::SCFDialect" + "scf::SCFDialect", + "tensor::TensorDialect" ]; } @@ -504,7 +505,7 @@ def TosaToStandard : Pass<"tosa-to-standard"> { let summary = "Lower TOSA to the Standard dialect"; - let dependentDialects = ["StandardOpsDialect"]; + let dependentDialects = ["StandardOpsDialect", "tensor::TensorDialect"]; let description = [{ Pass that converts TOSA operations to the equivalent operations using the operations in the Standard dialect. diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td --- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td +++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td @@ -33,7 +33,8 @@ }]; let cppNamespace = "::mlir::linalg"; let dependentDialects = [ - "AffineDialect", "StandardOpsDialect", "tensor::TensorDialect" + "AffineDialect", "memref::MemRefDialect", "StandardOpsDialect", + "tensor::TensorDialect" ]; let hasCanonicalizer = 1; let hasOperationAttrVerify = 1; diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.h b/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.h --- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.h +++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.h @@ -10,6 +10,7 @@ #define MLIR_DIALECT_LINALG_LINALGOPS_H_ #include "mlir/Dialect/Linalg/IR/LinalgTypes.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" #include "mlir/Dialect/Utils/StructuredOpsUtils.h" #include "mlir/IR/AffineExpr.h" diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgTypes.h b/mlir/include/mlir/Dialect/Linalg/IR/LinalgTypes.h --- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgTypes.h +++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgTypes.h @@ -10,6 +10,7 @@ #define MLIR_DIALECT_LINALG_LINALGTYPES_H_ #include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/Dialect.h" diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h --- a/mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h +++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h @@ -47,7 +47,7 @@ /// If hoistPaddingOnTensors is called with `nLoops` = 2 on the following IR. /// ``` /// scf.for (%i, %j, %k) -/// %st0 = subtensor f(%i, %k) : ... to tensor +/// %st0 = tensor.subtensor f(%i, %k) : ... to tensor /// %0 = linalg.pad_tensor %st0 low[0, 0] high[...] { /// ^bb0( ... ): /// linalg.yield %pad @@ -61,16 +61,17 @@ /// scf.for (%i) { /// %packed_init = linalg.init_tensor range(%j) : tensor /// %packed = scf.for (%k) iter_args(%p : %packed_init) { -/// %st0 = subtensor f(%i, %k) : ... to tensor +/// %st0 = tensor.subtensor f(%i, %k) : ... to tensor /// %0 = linalg.pad_tensor %st0 low[0, 0] high[...] { /// ^bb0( ... ): /// linalg.yield %pad /// } : tensor to tensor<4x8xf32> -/// %1 = subtensor_insert %0 ... : tensor<4x8xf32> to tensor +/// %1 = tensor.subtensor_insert %0 ... +/// : tensor<4x8xf32> to tensor /// scf.yield %1: tensor /// } -> tensor /// scf.for (%j, %k) { -/// %st0 = subtensor %packed [%k, 0, 0][1, 4, 8][1, 1, 1] : +/// %st0 = tensor.subtensor %packed [%k, 0, 0][1, 4, 8][1, 1, 1] : /// tensor to tensor<4x8xf32> /// compute(%st0) /// } diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h --- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h +++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h @@ -10,6 +10,7 @@ #define DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H_ #include "mlir/Dialect/Linalg/Utils/Utils.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/Utils.h" #include "mlir/Dialect/Vector/VectorOps.h" #include "mlir/IR/Identifier.h" diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h b/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h --- a/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h +++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h @@ -9,6 +9,7 @@ #ifndef MLIR_DIALECT_MEMREF_IR_MEMREF_H_ #define MLIR_DIALECT_MEMREF_IR_MEMREF_H_ +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/Dialect.h" #include "mlir/Interfaces/CallInterfaces.h" #include "mlir/Interfaces/CastInterfaces.h" diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefBase.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefBase.td --- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefBase.td +++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefBase.td @@ -19,6 +19,8 @@ manipulation ops, which are not strongly associated with any particular other dialect or domain abstraction. }]; + let dependentDialects = ["tensor::TensorDialect"]; + let hasConstantMaterializer = 1; } diff --git a/mlir/include/mlir/Dialect/Shape/IR/Shape.h b/mlir/include/mlir/Dialect/Shape/IR/Shape.h --- a/mlir/include/mlir/Dialect/Shape/IR/Shape.h +++ b/mlir/include/mlir/Dialect/Shape/IR/Shape.h @@ -14,6 +14,7 @@ #ifndef MLIR_SHAPE_IR_SHAPE_H #define MLIR_SHAPE_IR_SHAPE_H +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/Dialect.h" #include "mlir/IR/OpDefinition.h" diff --git a/mlir/include/mlir/Dialect/Shape/IR/ShapeBase.td b/mlir/include/mlir/Dialect/Shape/IR/ShapeBase.td --- a/mlir/include/mlir/Dialect/Shape/IR/ShapeBase.td +++ b/mlir/include/mlir/Dialect/Shape/IR/ShapeBase.td @@ -35,6 +35,7 @@ }]; let cppNamespace = "::mlir::shape"; + let dependentDialects = ["tensor::TensorDialect"]; let hasConstantMaterializer = 1; let hasOperationAttrVerify = 1; diff --git a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h --- a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h +++ b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h @@ -14,7 +14,6 @@ #ifndef MLIR_DIALECT_STANDARDOPS_IR_OPS_H #define MLIR_DIALECT_STANDARDOPS_IR_OPS_H -#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Dialect.h" @@ -24,7 +23,6 @@ #include "mlir/Interfaces/ControlFlowInterfaces.h" #include "mlir/Interfaces/SideEffectInterfaces.h" #include "mlir/Interfaces/VectorInterfaces.h" -#include "mlir/Interfaces/ViewLikeInterface.h" // Pull in all enum type definitions and utility function declarations. #include "mlir/Dialect/StandardOps/IR/OpsEnums.h.inc" @@ -35,12 +33,6 @@ class FuncOp; class OpBuilder; class PatternRewriter; - -/// Return the list of Range (i.e. offset, size, stride). Each Range -/// entry contains either the dynamic value or a ConstantIndexOp constructed -/// with `b` at location `loc`. -SmallVector getOrCreateRanges(OffsetSizeAndStrideOpInterface op, - OpBuilder &b, Location loc); } // namespace mlir #define GET_OP_CLASSES 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 @@ -21,15 +21,11 @@ include "mlir/Interfaces/ControlFlowInterfaces.td" include "mlir/Interfaces/SideEffectInterfaces.td" include "mlir/Interfaces/VectorInterfaces.td" -include "mlir/Interfaces/ViewLikeInterface.td" def StandardOps_Dialect : Dialect { let name = "std"; let cppNamespace = "::mlir"; let hasConstantMaterializer = 1; - // TODO: This dependency is needed to handle memref ops in the - // canonicalize pass and should be resolved. - let dependentDialects = ["memref::MemRefDialect"]; } // Base class for Standard dialect ops. @@ -1757,245 +1753,6 @@ let hasCanonicalizer = 1; } -//===----------------------------------------------------------------------===// -// SubTensorOp -//===----------------------------------------------------------------------===// - -def SubTensorOp : BaseOpWithOffsetSizesAndStrides< - StandardOps_Dialect, "subtensor", [NoSideEffect, AttrSizedOperandSegments, - OffsetSizeAndStrideOpInterface]> { - let summary = "subtensor operation"; - let description = [{ - The "subtensor" operation extract a tensor from another tensor as - specified by the operation's offsets, sizes and strides arguments. - - The subtensor operation supports the following arguments: - - * source: the "base" tensor from which to extract a subtensor. - * offsets: tensor-rank number of offsets into the "base" tensor from which - to extract the subtensor. - * sizes: tensor-rank number of sizes which specify the sizes of the result - tensor type. - * strides: tensor-rank number of strides specifying subsampling in each - dimension. - - The representation based on offsets, sizes and strides support a - partially-static specification via attributes specified through the - `static_offsets`, `static_sizes` and `static_strides` arguments. A special - sentinel value ShapedType::kDynamicSize and - ShapedType::kDynamicStrideOrOffset encodes that the corresponding entry has - a dynamic value. - - After buffer-allocation, the "subtensor" op is expected to lower into a - "subview" op. - - A subtensor operation may additionally reduce the rank of the resulting - tensor by removing dimensions that are statically known to be of size 1. - - Example: - - ``` - // Rank-reducing subtensor. - %1 = subtensor %0[0, 0, 0][1, 16, 4][1, 1, 1] : - tensor<8x16x4xf32> to tensor<16x4xf32> - %3 = subtensor %2[3, 4, 2][1, 6, 3][1, 1, 1] : - tensor<8x16x4xf32> to tensor<6x3xf32> - ``` - }]; - - let arguments = (ins - AnyRankedTensor:$source, - Variadic:$offsets, - Variadic:$sizes, - Variadic:$strides, - I64ArrayAttr:$static_offsets, - I64ArrayAttr:$static_sizes, - I64ArrayAttr:$static_strides - ); - let results = (outs AnyRankedTensor:$result); - - let assemblyFormat = [{ - $source `` - custom($offsets, $static_offsets) - custom($sizes, $static_sizes) - custom($strides, $static_strides) - attr-dict `:` type($source) `to` type($result) - }]; - - let builders = [ - // Build a SubTensorOp with mixed static and dynamic entries and inferred - // result type. - OpBuilder<(ins "Value":$source, "ArrayRef":$offsets, - "ArrayRef":$sizes, "ArrayRef":$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - // Build a SubTensorOp with mixed static and dynamic entries and custom - // result type. If the type passed is nullptr, it is inferred. - OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, - "ArrayRef":$offsets, "ArrayRef":$sizes, - "ArrayRef":$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - // Build a SubTensorOp with dynamic entries and custom result type. If the - // type passed is nullptr, it is inferred. - OpBuilder<(ins "Value":$source, "ValueRange":$offsets, - "ValueRange":$sizes, "ValueRange":$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - // Build a SubTensorOp with dynamic entries and inferred result type. - OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, - "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, - CArg<"ArrayRef", "{}">:$attrs)> - ]; - - let extraClassDeclaration = extraBaseClassDeclaration # [{ - /// Returns the type of the base tensor operand. - RankedTensorType getSourceType() { - return source().getType().cast(); - } - - /// The result of a subtensor is always a tensor. - RankedTensorType getType() { - return getResult().getType().cast(); - } - - /// A subtensor result type can be fully inferred from the source type and - /// the static representation of offsets, sizes and strides. Special - /// sentinels encode the dynamic case. - static Type inferResultType(RankedTensorType sourceRankedTensorType, - ArrayRef staticOffsets, - ArrayRef staticSizes, - ArrayRef staticStrides); - static Type inferResultType(RankedTensorType sourceRankedTensorType, - ArrayRef staticOffsets, - ArrayRef staticSizes, - ArrayRef staticStrides); - static Type inferRankReducedResultType(unsigned resultRank, - RankedTensorType sourceRankedTensorType, - ArrayRef staticOffsets, - ArrayRef staticSizes, - ArrayRef staticStrides); - static Type inferRankReducedResultType(unsigned resultRank, - RankedTensorType sourceRankedTensorType, - ArrayRef staticOffsets, - ArrayRef staticSizes, - ArrayRef staticStrides); - - /// Return the expected rank of each of the`static_offsets`, `static_sizes` - /// and `static_strides` attributes. - std::array getArrayAttrMaxRanks() { - unsigned rank = getSourceType().getRank(); - return {rank, rank, rank}; - } - - /// Return the number of leading operands before the `offsets`, `sizes` and - /// and `strides` operands. - static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; } - }]; - - let hasCanonicalizer = 1; - let hasFolder = 1; -} - -//===----------------------------------------------------------------------===// -// SubTensorInsertOp -//===----------------------------------------------------------------------===// - -def SubTensorInsertOp : BaseOpWithOffsetSizesAndStrides< - StandardOps_Dialect, "subtensor_insert", - [NoSideEffect, AttrSizedOperandSegments, OffsetSizeAndStrideOpInterface, - TypesMatchWith<"expected result type to match dest type", - "dest", "result", "$_self">]> { - let summary = "subtensor_insert operation"; - let description = [{ - The "subtensor_insert" operation insert a tensor `source` into another - tensor `dest` as specified by the operation's offsets, sizes and strides - arguments. - - It returns a copy of `dest` with the proper subtensor updated with the value - of `source`. - - The subtensor_insert operation has the encodes the following information: - - * source: the tensor that is inserted. - * dest: the tensor into which the source tensor is inserted. - * offsets: tensor-rank number of offsets into the "base" tensor from which - to extract the subtensor. - * sizes: tensor-rank number of sizes which specify the sizes of the result - tensor type. - * strides: tensor-rank number of strides that specify subsampling in each - dimension. - - The representation based on offsets, sizes and strides support a - partially-static specification via attributes specified through the - `static_offsets`, `static_sizes` and `static_strides` arguments. A special - sentinel value ShapedType::kDynamicSize and - ShapedType::kDynamicStrideOrOffset encodes that the corresponding entry has - a dynamic value. - - After buffer-allocation, the "subtensor_insert" op is expected to become - an in-place buffer update. - }]; - - let arguments = (ins - AnyRankedTensor:$source, - AnyRankedTensor:$dest, - Variadic:$offsets, - Variadic:$sizes, - Variadic:$strides, - I64ArrayAttr:$static_offsets, - I64ArrayAttr:$static_sizes, - I64ArrayAttr:$static_strides - ); - let results = (outs AnyRankedTensor:$result); - - let assemblyFormat = [{ - $source `into` $dest `` - custom($offsets, $static_offsets) - custom($sizes, $static_sizes) - custom($strides, $static_strides) - attr-dict `:` type($source) `into` type($dest) - }]; - - let verifier = ?; - - let builders = [ - // Build a SubTensorInsertOp with mixed static and dynamic entries. - OpBuilder<(ins "Value":$source, "Value":$dest, - "ArrayRef":$offsets, "ArrayRef":$sizes, - "ArrayRef":$strides, - CArg<"ArrayRef", "{}">:$attrs)>, - // Build a SubTensorInsertOp with dynamic entries. - OpBuilder<(ins "Value":$source, "Value":$dest, - "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, - CArg<"ArrayRef", "{}">:$attrs)> - ]; - - let extraClassDeclaration = extraBaseClassDeclaration # [{ - /// Returns the type of the base tensor operand. - RankedTensorType getSourceType() { - return source().getType().cast(); - } - - /// The result of a subtensor_insert is always a tensor. - RankedTensorType getType() { - return getResult().getType().cast(); - } - - /// Return the expected rank of each of the`static_offsets`, `static_sizes` - /// and `static_strides` attributes. - std::array getArrayAttrMaxRanks() { - unsigned rank = getType().getRank(); - return {rank, rank, rank}; - } - - /// Return the number of leading operands before the `offsets`, `sizes` and - /// and `strides` operands. - static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 2; } - }]; - - let hasCanonicalizer = 1; - let hasFolder = 1; -} - - //===----------------------------------------------------------------------===// // SwitchOp //===----------------------------------------------------------------------===// 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 @@ -16,6 +16,21 @@ #include "mlir/Interfaces/CastInterfaces.h" #include "mlir/Interfaces/ControlFlowInterfaces.h" #include "mlir/Interfaces/SideEffectInterfaces.h" +#include "mlir/Interfaces/ViewLikeInterface.h" + +//===----------------------------------------------------------------------===// +// Tensor Dialect Helpers +//===----------------------------------------------------------------------===// + +namespace mlir { + +/// Return the list of Range (i.e. offset, size, stride). Each Range +/// entry contains either the dynamic value or a ConstantIndexOp constructed +/// with `b` at location `loc`. +SmallVector getOrCreateRanges(OffsetSizeAndStrideOpInterface op, + OpBuilder &b, Location loc); + +} // namespace mlir //===----------------------------------------------------------------------===// // Tensor Dialect @@ -64,7 +79,6 @@ /// Performs folding of any operand of `op` if it comes from a tensor::CastOp /// that can be folded. LogicalResult foldTensorCast(Operation *op); - } // namespace tensor } // namespace mlir 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 @@ -13,6 +13,7 @@ include "mlir/Interfaces/CastInterfaces.td" include "mlir/Interfaces/ControlFlowInterfaces.td" include "mlir/Interfaces/SideEffectInterfaces.td" +include "mlir/Interfaces/ViewLikeInterface.td" class Tensor_Op traits = []> : Op { @@ -295,6 +296,244 @@ }]; } +//===----------------------------------------------------------------------===// +// SubTensorOp +//===----------------------------------------------------------------------===// + +def Tensor_SubTensorOp : BaseOpWithOffsetSizesAndStrides< + Tensor_Dialect, "subtensor", [NoSideEffect, AttrSizedOperandSegments, + OffsetSizeAndStrideOpInterface]> { + let summary = "subtensor operation"; + let description = [{ + The "subtensor" operation extract a tensor from another tensor as + specified by the operation's offsets, sizes and strides arguments. + + The subtensor operation supports the following arguments: + + * source: the "base" tensor from which to extract a subtensor. + * offsets: tensor-rank number of offsets into the "base" tensor from which + to extract the subtensor. + * sizes: tensor-rank number of sizes which specify the sizes of the result + tensor type. + * strides: tensor-rank number of strides specifying subsampling in each + dimension. + + The representation based on offsets, sizes and strides support a + partially-static specification via attributes specified through the + `static_offsets`, `static_sizes` and `static_strides` arguments. A special + sentinel value ShapedType::kDynamicSize and + ShapedType::kDynamicStrideOrOffset encodes that the corresponding entry has + a dynamic value. + + After buffer-allocation, the "subtensor" op is expected to lower into a + "subview" op. + + A subtensor operation may additionally reduce the rank of the resulting + tensor by removing dimensions that are statically known to be of size 1. + + Example: + + ``` + // Rank-reducing subtensor. + %1 = tensor.subtensor %0[0, 0, 0][1, 16, 4][1, 1, 1] : + tensor<8x16x4xf32> to tensor<16x4xf32> + %3 = tensor.subtensor %2[3, 4, 2][1, 6, 3][1, 1, 1] : + tensor<8x16x4xf32> to tensor<6x3xf32> + ``` + }]; + + let arguments = (ins + AnyRankedTensor:$source, + Variadic:$offsets, + Variadic:$sizes, + Variadic:$strides, + I64ArrayAttr:$static_offsets, + I64ArrayAttr:$static_sizes, + I64ArrayAttr:$static_strides + ); + let results = (outs AnyRankedTensor:$result); + + let assemblyFormat = [{ + $source `` + custom($offsets, $static_offsets) + custom($sizes, $static_sizes) + custom($strides, $static_strides) + attr-dict `:` type($source) `to` type($result) + }]; + + let builders = [ + // Build a SubTensorOp with mixed static and dynamic entries and inferred + // result type. + OpBuilder<(ins "Value":$source, "ArrayRef":$offsets, + "ArrayRef":$sizes, "ArrayRef":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, + // Build a SubTensorOp with mixed static and dynamic entries and custom + // result type. If the type passed is nullptr, it is inferred. + OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, + "ArrayRef":$offsets, "ArrayRef":$sizes, + "ArrayRef":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, + // Build a SubTensorOp with dynamic entries and custom result type. If the + // type passed is nullptr, it is inferred. + OpBuilder<(ins "Value":$source, "ValueRange":$offsets, + "ValueRange":$sizes, "ValueRange":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, + // Build a SubTensorOp with dynamic entries and inferred result type. + OpBuilder<(ins "RankedTensorType":$resultType, "Value":$source, + "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, + CArg<"ArrayRef", "{}">:$attrs)> + ]; + + let extraClassDeclaration = extraBaseClassDeclaration # [{ + /// Returns the type of the base tensor operand. + RankedTensorType getSourceType() { + return source().getType().cast(); + } + + /// The result of a subtensor is always a tensor. + RankedTensorType getType() { + return getResult().getType().cast(); + } + + /// A subtensor result type can be fully inferred from the source type and + /// the static representation of offsets, sizes and strides. Special + /// sentinels encode the dynamic case. + static Type inferResultType(RankedTensorType sourceRankedTensorType, + ArrayRef staticOffsets, + ArrayRef staticSizes, + ArrayRef staticStrides); + static Type inferResultType(RankedTensorType sourceRankedTensorType, + ArrayRef staticOffsets, + ArrayRef staticSizes, + ArrayRef staticStrides); + static Type inferRankReducedResultType(unsigned resultRank, + RankedTensorType sourceRankedTensorType, + ArrayRef staticOffsets, + ArrayRef staticSizes, + ArrayRef staticStrides); + static Type inferRankReducedResultType(unsigned resultRank, + RankedTensorType sourceRankedTensorType, + ArrayRef staticOffsets, + ArrayRef staticSizes, + ArrayRef staticStrides); + + /// Return the expected rank of each of the`static_offsets`, `static_sizes` + /// and `static_strides` attributes. + std::array getArrayAttrMaxRanks() { + unsigned rank = getSourceType().getRank(); + return {rank, rank, rank}; + } + + /// Return the number of leading operands before the `offsets`, `sizes` and + /// and `strides` operands. + static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; } + }]; + + let hasCanonicalizer = 1; + let hasFolder = 1; +} + +//===----------------------------------------------------------------------===// +// SubTensorInsertOp +//===----------------------------------------------------------------------===// + +def Tensor_SubTensorInsertOp : BaseOpWithOffsetSizesAndStrides< + Tensor_Dialect, "subtensor_insert", + [NoSideEffect, AttrSizedOperandSegments, OffsetSizeAndStrideOpInterface, + TypesMatchWith<"expected result type to match dest type", + "dest", "result", "$_self">]> { + let summary = "subtensor_insert operation"; + let description = [{ + The "subtensor_insert" operation insert a tensor `source` into another + tensor `dest` as specified by the operation's offsets, sizes and strides + arguments. + + It returns a copy of `dest` with the proper subtensor updated with the value + of `source`. + + The subtensor_insert operation has the encodes the following information: + + * source: the tensor that is inserted. + * dest: the tensor into which the source tensor is inserted. + * offsets: tensor-rank number of offsets into the "base" tensor from which + to extract the subtensor. + * sizes: tensor-rank number of sizes which specify the sizes of the result + tensor type. + * strides: tensor-rank number of strides that specify subsampling in each + dimension. + + The representation based on offsets, sizes and strides support a + partially-static specification via attributes specified through the + `static_offsets`, `static_sizes` and `static_strides` arguments. A special + sentinel value ShapedType::kDynamicSize and + ShapedType::kDynamicStrideOrOffset encodes that the corresponding entry has + a dynamic value. + + After buffer-allocation, the "subtensor_insert" op is expected to become + an in-place buffer update. + }]; + + let arguments = (ins + AnyRankedTensor:$source, + AnyRankedTensor:$dest, + Variadic:$offsets, + Variadic:$sizes, + Variadic:$strides, + I64ArrayAttr:$static_offsets, + I64ArrayAttr:$static_sizes, + I64ArrayAttr:$static_strides + ); + let results = (outs AnyRankedTensor:$result); + + let assemblyFormat = [{ + $source `into` $dest `` + custom($offsets, $static_offsets) + custom($sizes, $static_sizes) + custom($strides, $static_strides) + attr-dict `:` type($source) `into` type($dest) + }]; + + let verifier = ?; + + let builders = [ + // Build a SubTensorInsertOp with mixed static and dynamic entries. + OpBuilder<(ins "Value":$source, "Value":$dest, + "ArrayRef":$offsets, "ArrayRef":$sizes, + "ArrayRef":$strides, + CArg<"ArrayRef", "{}">:$attrs)>, + // Build a SubTensorInsertOp with dynamic entries. + OpBuilder<(ins "Value":$source, "Value":$dest, + "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides, + CArg<"ArrayRef", "{}">:$attrs)> + ]; + + let extraClassDeclaration = extraBaseClassDeclaration # [{ + /// Returns the type of the base tensor operand. + RankedTensorType getSourceType() { + return source().getType().cast(); + } + + /// The result of a subtensor_insert is always a tensor. + RankedTensorType getType() { + return getResult().getType().cast(); + } + + /// Return the expected rank of each of the`static_offsets`, `static_sizes` + /// and `static_strides` attributes. + std::array getArrayAttrMaxRanks() { + unsigned rank = getType().getRank(); + return {rank, rank, rank}; + } + + /// Return the number of leading operands before the `offsets`, `sizes` and + /// and `strides` operands. + static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 2; } + }]; + + let hasCanonicalizer = 1; + let hasFolder = 1; +} + //===----------------------------------------------------------------------===// // YieldOp //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp --- a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp +++ b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp @@ -18,6 +18,7 @@ #include "mlir/Dialect/GPU/Passes.h" #include "mlir/Dialect/LLVMIR/NVVMDialect.h" #include "mlir/Dialect/Math/IR/Math.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/IR/BlockAndValueMapping.h" #include "mlir/Transforms/DialectConversion.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h" 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 @@ -627,10 +627,11 @@ Value index = b.create(indexIsNegative, add, originalIndex); Value one = b.create(1); - Value head = b.create(transformed.operand(), zero, index, one); + Value head = + b.create(transformed.operand(), zero, index, one); Value tailSize = b.create(rank, index); - Value tail = - b.create(transformed.operand(), index, tailSize, one); + Value tail = b.create(transformed.operand(), index, + tailSize, one); rewriter.replaceOp(op, {head, tail}); return success(); } diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp --- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp +++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp @@ -1741,8 +1741,8 @@ for (auto arg : args) { sizes[axis] = rewriter.create(loc, arg, axisValue); - result = rewriter.create(loc, arg, result, offsets, - sizes, strides); + result = rewriter.create( + loc, arg, result, offsets, sizes, strides); offsets[axis] = rewriter.create(loc, offsets[axis], sizes[axis]); } rewriter.replaceOp(op, result); diff --git a/mlir/lib/Conversion/TosaToStandard/TosaToStandard.cpp b/mlir/lib/Conversion/TosaToStandard/TosaToStandard.cpp --- a/mlir/lib/Conversion/TosaToStandard/TosaToStandard.cpp +++ b/mlir/lib/Conversion/TosaToStandard/TosaToStandard.cpp @@ -12,6 +12,7 @@ #include "mlir/Conversion/TosaToStandard/TosaToStandard.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Tosa/IR/TosaOps.h" #include "mlir/IR/PatternMatch.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h" @@ -42,7 +43,7 @@ SmallVector strides; strides.resize(sliceOp.getType().template cast().getRank(), 1); - rewriter.replaceOpWithNewOp( + rewriter.replaceOpWithNewOp( sliceOp, sliceOp.getType(), input, ValueRange({}), ValueRange({}), ValueRange({}), sliceOp.start(), sliceOp.size(), rewriter.getI64ArrayAttr(strides)); diff --git a/mlir/lib/Conversion/TosaToStandard/TosaToStandardPass.cpp b/mlir/lib/Conversion/TosaToStandard/TosaToStandardPass.cpp --- a/mlir/lib/Conversion/TosaToStandard/TosaToStandardPass.cpp +++ b/mlir/lib/Conversion/TosaToStandard/TosaToStandardPass.cpp @@ -35,6 +35,7 @@ target.addIllegalOp(); target.addIllegalOp(); target.addLegalDialect(); + target.addLegalDialect(); mlir::tosa::populateTosaToStandardConversionPatterns(&patterns); if (failed(applyPartialConversion(getOperation(), target, diff --git a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp --- a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp +++ b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp @@ -17,6 +17,7 @@ #include "../PassDetail.h" #include "mlir/Analysis/SliceAnalysis.h" #include "mlir/Dialect/GPU/GPUDialect.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/Utils/StructuredOpsUtils.h" #include "mlir/Dialect/Vector/VectorOps.h" #include "mlir/Dialect/Vector/VectorUtils.h" diff --git a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp --- a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp +++ b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp @@ -17,6 +17,7 @@ #include "../PassDetail.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/Utils.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/SCF.h" #include "mlir/Dialect/Vector/VectorOps.h" #include "mlir/Dialect/Vector/VectorUtils.h" diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineDataCopyGeneration.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineDataCopyGeneration.cpp --- a/mlir/lib/Dialect/Affine/Transforms/AffineDataCopyGeneration.cpp +++ b/mlir/lib/Dialect/Affine/Transforms/AffineDataCopyGeneration.cpp @@ -23,6 +23,7 @@ #include "mlir/Analysis/Utils.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/Passes.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h" #include "mlir/Transforms/LoopUtils.h" diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp --- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp +++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp @@ -16,6 +16,7 @@ #include "mlir/Dialect/Linalg/IR/LinalgTypes.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/AffineExprVisitor.h" #include "mlir/IR/Matchers.h" #include "mlir/IR/OpImplementation.h" @@ -749,10 +750,11 @@ /// subtensor of this is also needed only for its shape. The result can be /// replaced by a new init_tensor operation of the same size as the subtensor /// op. -struct FoldInitTensorWithSubTensorOp : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; +struct FoldInitTensorWithSubTensorOp + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(SubTensorOp subtensorOp, + LogicalResult matchAndRewrite(tensor::SubTensorOp subtensorOp, PatternRewriter &rewriter) const override { if (!subtensorOp.source().getDefiningOp()) return failure(); diff --git a/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp @@ -15,6 +15,7 @@ #include "mlir/Dialect/Math/IR/Math.h" #include "mlir/Dialect/StandardOps/Transforms/Passes.h" #include "mlir/Dialect/StandardOps/Utils/Utils.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Vector/VectorOps.h" #include "mlir/IR/BuiltinDialect.h" #include "mlir/IR/Operation.h" @@ -242,14 +243,14 @@ /// /// This pattern is arguable a std pattern once linalg::CopyOp becomes /// std::CopyOp. -class SubTensorOpConverter : public OpConversionPattern { +class SubTensorOpConverter : public OpConversionPattern { public: - using OpConversionPattern::OpConversionPattern; + using OpConversionPattern::OpConversionPattern; LogicalResult - matchAndRewrite(SubTensorOp op, ArrayRef operands, + matchAndRewrite(tensor::SubTensorOp op, ArrayRef operands, ConversionPatternRewriter &rewriter) const final { - SubTensorOpAdaptor adaptor(operands, op->getAttrDictionary()); + tensor::SubTensorOpAdaptor adaptor(operands, op->getAttrDictionary()); Value sourceMemref = adaptor.source(); assert(sourceMemref.getType().isa()); @@ -282,14 +283,14 @@ /// This pattern is arguable a std pattern once linalg::CopyOp becomes /// std::CopyOp. class SubTensorInsertOpConverter - : public OpConversionPattern { + : public OpConversionPattern { public: - using OpConversionPattern::OpConversionPattern; + using OpConversionPattern::OpConversionPattern; LogicalResult - matchAndRewrite(SubTensorInsertOp op, ArrayRef operands, + matchAndRewrite(tensor::SubTensorInsertOp op, ArrayRef operands, ConversionPatternRewriter &rewriter) const final { - SubTensorInsertOpAdaptor adaptor(operands, op->getAttrDictionary()); + tensor::SubTensorInsertOpAdaptor adaptor(operands, op->getAttrDictionary()); Value sourceMemRef = adaptor.source(); assert(sourceMemRef.getType().isa()); @@ -323,7 +324,8 @@ // Mark all Standard operations legal. target.addLegalDialect(); - target.addIllegalOp(); + target.addIllegalOp(); // Mark all Linalg operations illegal as long as they work on tensors. auto isLegalOperation = [&](Operation *op) { diff --git a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp @@ -82,6 +82,7 @@ #include "mlir/Dialect/Linalg/IR/LinalgOps.h" #include "mlir/Dialect/Linalg/Passes.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Vector/VectorOps.h" #include "mlir/IR/Operation.h" #include "mlir/Pass/Pass.h" @@ -141,7 +142,8 @@ /// Return the OpResult that matches an operand. /// Return null if no such result exists. -OpResult getMatchingOpResult(SubTensorInsertOp op, OpOperand &opOperand) { +OpResult getMatchingOpResult(tensor::SubTensorInsertOp op, + OpOperand &opOperand) { if (opOperand.get() != op.dest()) return OpResult(); return op->getResult(0); @@ -158,11 +160,11 @@ // Ops that perform destructive updates on operand(s) to produce // result(s). .Case( [&](auto op) { return getMatchingOpResult(op, opOperand); }) // Other ops. - .Case([&](auto op) { return OpResult(); }) + .Case([&](auto op) { return OpResult(); }) .Default([&](Operation *op) { return OpResult(); }); // clang-format on } @@ -563,7 +565,7 @@ /// whether or not it is marked inplaceable. /// Note that `getMatchingOpResult` on a SubTensorOp always returns null. /// As consequence a SubTensorOp always alloc + copy when taken in isolation. -static LogicalResult bufferize(OpBuilder &b, SubTensorOp subTensorOp, +static LogicalResult bufferize(OpBuilder &b, tensor::SubTensorOp subTensorOp, BlockAndValueMapping &bvm) { LLVM_DEBUG(DBGS() << "bufferize: " << *subTensorOp << "\n"); @@ -609,7 +611,7 @@ } static LogicalResult bufferize(OpBuilder &b, - SubTensorInsertOp subTensorInsertOp, + tensor::SubTensorInsertOp subTensorInsertOp, BlockAndValueMapping &bvm) { LLVM_DEBUG(DBGS() << "bufferize: " << *subTensorInsertOp << "\n"); @@ -822,7 +824,7 @@ static bool subTensorUseBufferizesToReadOnly(OpOperand &operand) { assert(operand.get().getDefiningOp() && "expected subtensor op"); if (auto subTensorInsertOp = - dyn_cast(operand.getOwner())) { + dyn_cast(operand.getOwner())) { return operand.getOperandNumber() == 0 /* source of the subTensorInsert*/ && // If the subTensorInsertOp is not inplace, there is no possible // internal aliasing with subTensorOp, which is inplaceable. @@ -848,7 +850,7 @@ /// SubTensorOp introduces potential aliasing and a combination of things need /// to occur to determine whether it is inplaceable. -static void analyzeInPlaceSubTensor(SubTensorOp subTensorOp, +static void analyzeInPlaceSubTensor(tensor::SubTensorOp subTensorOp, const DominanceInfo &domInfo) { // Case 1: // a. All uses are known to bufferize to readonly buffers. @@ -871,7 +873,7 @@ funcOp.walk([&](Operation *op) { // Skip SubTensorOp in a first pass. - if (auto subTensorOp = dyn_cast(op)) + if (auto subTensorOp = dyn_cast(op)) return analyzeInPlaceSubTensor(subTensorOp, domInfo); // All other ops are checked for `isBufferizableInPlace`. @@ -904,8 +906,8 @@ .Case( [&](auto op) { return bufferize(b, op, bvm); }) // clang-format on diff --git a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp @@ -18,6 +18,7 @@ #include "mlir/Dialect/Linalg/Passes.h" #include "mlir/Dialect/Linalg/Transforms/Transforms.h" #include "mlir/Dialect/Linalg/Utils/Utils.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/AffineExpr.h" #include "mlir/IR/AffineMap.h" #include "mlir/IR/BuiltinTypes.h" @@ -487,10 +488,11 @@ namespace { /// Convert `subtensor` operations to rank-reduced versions. -struct UseRankReducedSubTensorOp : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; +struct UseRankReducedSubTensorOp + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(SubTensorOp subTensorOp, + LogicalResult matchAndRewrite(tensor::SubTensorOp subTensorOp, PatternRewriter &rewriter) const override { RankedTensorType resultType = subTensorOp.getType(); SmallVector offsets = subTensorOp.getMixedOffsets(); @@ -501,13 +503,13 @@ reassociation->size() == static_cast(resultType.getRank())) return failure(); auto rankReducedType = - SubTensorOp::inferRankReducedResultType(reassociation->size(), - subTensorOp.getSourceType(), - offsets, sizes, strides) + tensor::SubTensorOp::inferRankReducedResultType( + reassociation->size(), subTensorOp.getSourceType(), offsets, sizes, + strides) .cast(); Location loc = subTensorOp.getLoc(); - Value newSubTensor = rewriter.create( + Value newSubTensor = rewriter.create( loc, rankReducedType, subTensorOp.source(), offsets, sizes, strides); rewriter.replaceOpWithNewOp( subTensorOp, resultType, newSubTensor, *reassociation); @@ -517,10 +519,10 @@ /// Convert `subtensor_insert` operations to rank-reduced versions. struct UseRankReducedSubTensorInsertOp - : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(SubTensorInsertOp insertOp, + LogicalResult matchAndRewrite(tensor::SubTensorInsertOp insertOp, PatternRewriter &rewriter) const override { RankedTensorType sourceType = insertOp.getSourceType(); SmallVector offsets = insertOp.getMixedOffsets(); @@ -533,7 +535,7 @@ Location loc = insertOp.getLoc(); auto reshapedSource = rewriter.create( loc, insertOp.source(), *reassociation); - rewriter.replaceOpWithNewOp( + rewriter.replaceOpWithNewOp( insertOp, reshapedSource, insertOp.dest(), insertOp.getMixedOffsets(), insertOp.getMixedSizes(), insertOp.getMixedStrides()); return success(); diff --git a/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp @@ -78,8 +78,9 @@ // `ViewInterface`. The interface needs a `getOrCreateRanges` method which // currently returns a `linalg.range`. The fix here is to move this op to // `std` dialect and add the method to `ViewInterface`. - if (fromSubViewOpOnly && !isa_and_nonnull( - opOperand->get().getDefiningOp())) + if (fromSubViewOpOnly && + !isa_and_nonnull( + opOperand->get().getDefiningOp())) continue; AffineMap map = op.getTiedIndexingMap(opOperand); @@ -221,7 +222,7 @@ SmallVector staticSizesVector(rank, ShapedType::kDynamicSize); SmallVector staticStridesVector( rank, ShapedType::kDynamicStrideOrOffset); - resultTypes.push_back(SubTensorOp::inferResultType( + resultTypes.push_back(tensor::SubTensorOp::inferResultType( t.cast(), staticOffsetsVector, staticSizesVector, staticStridesVector)); } @@ -258,7 +259,7 @@ Operation *shapeProducingOp = shapedOperand.getDefiningOp(); if (auto subViewOp = dyn_cast(shapeProducingOp)) return subViewOp.getOrCreateRanges(b, loc)[dim]; - if (auto subTensorOp = dyn_cast(shapeProducingOp)) + if (auto subTensorOp = dyn_cast(shapeProducingOp)) return subTensorOp.getOrCreateRanges(b, loc)[dim]; llvm_unreachable("SubviewOp or SubTensorOp expected"); } @@ -473,7 +474,7 @@ opResult = tensor.cast(); return; } - if (auto subTensorOp = tensor.getDefiningOp()) { + if (auto subTensorOp = tensor.getDefiningOp()) { tensor = subTensorOp.source(); continue; } @@ -513,7 +514,7 @@ Value inputTensor = consumerOpOperand.get(); // Must be a subtensor to guarantee there are loops we can fuse into. - auto subTensor = inputTensor.getDefiningOp(); + auto subTensor = inputTensor.getDefiningOp(); if (!subTensor) { LLVM_DEBUG(llvm::dbgs() << "\nNot fusable, not a subtensor: " << inputTensor); diff --git a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp @@ -19,6 +19,7 @@ #include "mlir/Dialect/SCF/SCF.h" #include "mlir/Dialect/SCF/Utils.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Vector/VectorOps.h" #include "mlir/Dialect/Vector/VectorUtils.h" #include "mlir/IR/BuiltinOps.h" @@ -42,13 +43,13 @@ /// instructions that need to be hoisted too. struct HoistableWrite { vector::TransferWriteOp transferWriteOp; - SubTensorInsertOp subTensorInsertOp; + tensor::SubTensorInsertOp subTensorInsertOp; }; /// Represents a unit of hoistable TransferReadOp. This may comprise other /// instructions that need to be hoisted too. struct HoistableRead { vector::TransferReadOp transferReadOp; - SubTensorOp subTensorOp; + tensor::SubTensorOp subTensorOp; }; } // namespace @@ -71,7 +72,8 @@ } /// Return true is all offsets, sizes and strides are equal. -static bool sameOffsetsSizesAndStrides(SubTensorOp s, SubTensorInsertOp si) { +static bool sameOffsetsSizesAndStrides(tensor::SubTensorOp s, + tensor::SubTensorInsertOp si) { if (s.static_offsets().size() != si.static_offsets().size()) return false; if (s.static_sizes().size() != si.static_sizes().size()) @@ -109,10 +111,10 @@ // If HoistableWrite involves a SubTensorInsertOp, we need to find a // matching SubTensorOp. - SubTensorOp subTensorOp; + tensor::SubTensorOp subTensorOp; Operation *maybeTransferReadUser = user; if (write.subTensorInsertOp) { - subTensorOp = dyn_cast(user); + subTensorOp = dyn_cast(user); if (!subTensorOp || subTensorOp.getResult().getType() != write.subTensorInsertOp.source().getType()) continue; @@ -174,7 +176,7 @@ // Consider all transitive uses through a subtensor / subtensor_insert. // TODO: atm we just bail because a stronger analysis is needed for these // cases. - if (isa(user)) + if (isa(user)) return true; // Consider all transitive uses through a vector.transfer_write. if (auto writeUser = dyn_cast(user)) { @@ -229,7 +231,7 @@ return HoistableWrite{write, nullptr}; } - if (auto subTensorInsertOp = v.getDefiningOp()) { + if (auto subTensorInsertOp = v.getDefiningOp()) { // Inserted subTensor must come from vector.transfer_write. auto write = subTensorInsertOp.source().getDefiningOp(); @@ -244,7 +246,8 @@ // Indexing inserted into must not depend on `forOp`. for (Value operand : subTensorInsertOp->getOperands().drop_front( - SubTensorInsertOp::getOffsetSizeAndStrideStartOperandIndex())) + tensor::SubTensorInsertOp:: + getOffsetSizeAndStrideStartOperandIndex())) if (!forOp.isDefinedOutsideOfLoop(operand)) return HoistableWrite(); @@ -801,7 +804,7 @@ for (Operation *op : backwardSlice) { // Specifically sit out in the subtenso(packedTensor) case: this is the // piece we seek to replace. - if (auto subTensor = dyn_cast(op)) + if (auto subTensor = dyn_cast(op)) if (bvm.lookupOrDefault(subTensor.source()) == packedTensor) continue; auto effects = dyn_cast(op); @@ -855,9 +858,9 @@ // strides = [1 .. 1]. SmallVector strides(nLoops + paddedRank, b.getIndexAttr(1)); - Value inserted = - b.create(loc, bvm.lookup(padTensorOp.result()), - packedTensor, offsets, sizes, strides); + Value inserted = b.create( + loc, bvm.lookup(padTensorOp.result()), packedTensor, offsets, sizes, + strides); // Stack step 3. iteratively pop the stack and propagate the yield. Value valueToYield = inserted; @@ -888,8 +891,8 @@ packedTensor = scf::getForInductionVarOwner(clonedLoopIvs.front())->getResult(0); padTensorOp.replaceAllUsesWith( - b.create(loc, padTensorOp.getResultType(), packedTensor, - offsets, sizes, strides) + b.create(loc, padTensorOp.getResultType(), + packedTensor, offsets, sizes, strides) ->getResult(0)); Operation *toErase = padTensorOp; diff --git a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp @@ -260,8 +260,8 @@ // TODO: use an interface/adaptor to avoid leaking position in // `tiledOperands`. Value outputTensor = tiledOperands[opOperand->getOperandNumber()]; - if (auto subtensor = outputTensor.getDefiningOp()) { - tensorResults.push_back(b.create( + if (auto subtensor = outputTensor.getDefiningOp()) { + tensorResults.push_back(b.create( loc, subtensor.source().getType(), res->getResult(resultIdx), subtensor.source(), subtensor.offsets(), subtensor.sizes(), subtensor.strides(), subtensor.static_offsets(), @@ -406,7 +406,7 @@ scf::ForOp::getCanonicalizationPatterns(patterns, ctx); scf::ParallelOp::getCanonicalizationPatterns(patterns, ctx); ConstantIndexOp::getCanonicalizationPatterns(patterns, ctx); - SubTensorOp::getCanonicalizationPatterns(patterns, ctx); + tensor::SubTensorOp::getCanonicalizationPatterns(patterns, ctx); memref::SubViewOp::getCanonicalizationPatterns(patterns, ctx); tensor::CastOp::getCanonicalizationPatterns(patterns, ctx); memref::ViewOp::getCanonicalizationPatterns(patterns, ctx); diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp @@ -16,6 +16,7 @@ #include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h" #include "mlir/Dialect/Linalg/IR/LinalgOps.h" #include "mlir/Dialect/Linalg/Utils/Utils.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Utils/StructuredOpsUtils.h" #include "mlir/Dialect/Vector/VectorOps.h" #include "mlir/IR/AffineExpr.h" @@ -128,7 +129,7 @@ // Already static shape, no need to pad. if (llvm::none_of(opToPad.getShape(opOperand), ShapedType::isDynamic)) return success(); - auto subtensor = opOperand->get().getDefiningOp(); + auto subtensor = opOperand->get().getDefiningOp(); // Not a subtensor, cannot construct a static bounding box. if (!subtensor) return failure(); @@ -211,7 +212,7 @@ return dimOp.getResult(); })); SmallVector strides(rank, rewriter.getIndexAttr(1)); - paddedSubviewResults.push_back(rewriter.create( + paddedSubviewResults.push_back(rewriter.create( loc, std::get<1>(it), offsets, sizes, strides)); } // Replace the transient `opToPad` locally, except for uses that we just diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp @@ -15,6 +15,7 @@ #include "mlir/Dialect/Linalg/IR/LinalgOps.h" #include "mlir/Dialect/Linalg/Transforms/Transforms.h" #include "mlir/Dialect/Linalg/Utils/Utils.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Utils/StructuredOpsUtils.h" #include "mlir/Dialect/Vector/VectorOps.h" #include "mlir/IR/AffineExpr.h" @@ -740,7 +741,7 @@ // Strides of SubTensorInsertOp are all 1. SmallVector strides(sourceType.getRank(), rewriter.getIndexAttr(1)); - rewriter.replaceOpWithNewOp( + rewriter.replaceOpWithNewOp( padOp, padOp.source(), fill, padOp.getMixedLowPad(), srcSizes, strides); return success(); @@ -913,16 +914,18 @@ /// write. In such cases, the TransferWriteOp can write to the non-padded tensor /// value and apply out-of-bounds masking. E.g.: /// ``` -/// %0 = subtensor ...[...] [%s0, %s1] [1, 1] : tensor<...> to tensor +/// %0 = tensor.subtensor ...[...] [%s0, %s1] [1, 1] +/// : tensor<...> to tensor /// %1 = linalg.pad_tensor %0 ... : tensor to tensor<17x5xf32> /// %2 = vector.transfer_write %vec, %1[...] /// : vector<17x5xf32>, tensor<17x5xf32> -/// %r = subtensor %2[0, 0] [%s0, %s1] [1, 1] +/// %r = tensor.subtensor %2[0, 0] [%s0, %s1] [1, 1] /// : tensor<17x5xf32> to tensor /// ``` /// is rewritten to: /// ``` -/// %0 = subtensor ...[...] [%s0, %s1] [1, 1] : tensor<...> to tensor +/// %0 = tensor.subtensor ...[...] [%s0, %s1] [1, 1] +/// : tensor<...> to tensor /// %r = vector.transfer_write %vec, %0[...] : vector<17x5xf32>, tensor /// ``` /// Note: It is important that the SubTensorOp %r resizes the result of the @@ -949,7 +952,7 @@ if (!padValue) return failure(); // TransferWriteOp result must be directly consumed by a SubTensorOp. if (!xferOp->hasOneUse()) return failure(); - auto trimPadding = dyn_cast(*xferOp->user_begin()); + auto trimPadding = dyn_cast(*xferOp->user_begin()); if (!trimPadding) return failure(); // Only static zero offsets supported when trimming padding. if (!trimPadding.hasZeroOffset()) return failure(); @@ -976,7 +979,8 @@ /// This is a conservative analysis. In case equal tensor sizes cannot be /// proven statically, this analysis returns `false` even though the tensor /// sizes may turn out to be equal at runtime. - bool hasSameTensorSize(Value beforePadding, SubTensorOp afterTrimming) const { + bool hasSameTensorSize(Value beforePadding, + tensor::SubTensorOp afterTrimming) const { // If the input to PadTensorOp is a CastOp, try with with both CastOp result // and CastOp operand. if (auto castOp = beforePadding.getDefiningOp()) @@ -1005,7 +1009,7 @@ // is when `beforePadding` is a SubTensorOp (or a cast thereof). // Apart from CastOp, only SubTensorOp is supported. - auto beforeSubtensor = beforePadding.getDefiningOp(); + auto beforeSubtensor = beforePadding.getDefiningOp(); if (!beforeSubtensor) return false; assert(static_cast(t1.getRank()) @@ -1045,7 +1049,8 @@ /// Rewrite use of PadTensorOp result in SubtensorInsertOp. E.g.: /// ``` /// %0 = linalg.pad_tensor %src ... : tensor to tensor<17x5xf32> -/// %r = subtensor_insert %0 into %dest[%a, %b, 0, 0] [1, 1, 17, 5] [1, 1, 1, 1] +/// %r = tensor.subtensor_insert %0 +/// into %dest[%a, %b, 0, 0] [1, 1, 17, 5] [1, 1, 1, 1] /// : tensor<17x5xf32> into tensor /// ``` /// is rewritten to: @@ -1064,12 +1069,12 @@ /// - Only unit strides in `insertOp`. /// - Single, scalar padding value. struct PadTensorOpVectorizationWithSubTensorInsertPattern - : public VectorizePadTensorOpUserPattern { - using VectorizePadTensorOpUserPattern - ::VectorizePadTensorOpUserPattern; + : public VectorizePadTensorOpUserPattern { + using VectorizePadTensorOpUserPattern< + tensor::SubTensorInsertOp>::VectorizePadTensorOpUserPattern; LogicalResult rewriteUser(PatternRewriter &rewriter, PadTensorOp padOp, - SubTensorInsertOp insertOp) const override { + tensor::SubTensorInsertOp insertOp) const override { // Low padding must be static 0. if (!padOp.hasZeroLowPad()) return failure(); // Only unit stride supported. diff --git a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp --- a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp +++ b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp @@ -15,9 +15,11 @@ #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Linalg/IR/LinalgOps.h" #include "mlir/Dialect/Linalg/IR/LinalgTypes.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/SCF.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" #include "mlir/Dialect/StandardOps/Utils/Utils.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/AffineExpr.h" #include "mlir/IR/AffineExprVisitor.h" #include "mlir/IR/AffineMap.h" @@ -618,8 +620,8 @@ tiledShapes.push_back( b.create(loc, shapedOp, offsets, sizes, strides)); else - tiledShapes.push_back( - b.create(loc, shapedOp, offsets, sizes, strides)); + tiledShapes.push_back(b.create( + loc, shapedOp, offsets, sizes, strides)); } return tiledShapes; diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp --- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp +++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp @@ -717,7 +717,7 @@ // The size at the given index is now known to be a dynamic size. unsigned unsignedIndex = index.getValue().getZExtValue(); - if (auto subtensor = dyn_cast_or_null(definingOp)) { + if (auto subtensor = dyn_cast_or_null(definingOp)) { assert(subtensor.isDynamicSize(unsignedIndex) && "Expected dynamic subtensor size"); return subtensor.getDynamicSize(unsignedIndex); diff --git a/mlir/lib/Dialect/SCF/Transforms/Bufferize.cpp b/mlir/lib/Dialect/SCF/Transforms/Bufferize.cpp --- a/mlir/lib/Dialect/SCF/Transforms/Bufferize.cpp +++ b/mlir/lib/Dialect/SCF/Transforms/Bufferize.cpp @@ -8,6 +8,7 @@ #include "mlir/Transforms/Bufferize.h" #include "PassDetail.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/Passes.h" #include "mlir/Dialect/SCF/SCF.h" #include "mlir/Dialect/SCF/Transforms.h" diff --git a/mlir/lib/Dialect/Shape/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Shape/Transforms/Bufferize.cpp --- a/mlir/lib/Dialect/Shape/Transforms/Bufferize.cpp +++ b/mlir/lib/Dialect/Shape/Transforms/Bufferize.cpp @@ -8,6 +8,7 @@ #include "mlir/Transforms/Bufferize.h" #include "PassDetail.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/Shape/Transforms/Passes.h" #include "mlir/Pass/Pass.h" diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp --- a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp @@ -43,6 +43,7 @@ #include "mlir/Dialect/Linalg/IR/LinalgOps.h" #include "mlir/Dialect/Linalg/Utils/Utils.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/SCF.h" #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h" #include "mlir/Dialect/SparseTensor/Transforms/Passes.h" 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 @@ -10,7 +10,6 @@ #include "mlir/Dialect/CommonFolders.h" #include "mlir/Dialect/StandardOps/Utils/Utils.h" -#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/AffineExpr.h" #include "mlir/IR/AffineMap.h" #include "mlir/IR/BlockAndValueMapping.h" @@ -34,32 +33,6 @@ using namespace mlir; -/// Helper function to dispatch an OpFoldResult into either the `dynamicVec` if -/// it is a Value or into `staticVec` if it is an IntegerAttr. -/// In the case of a Value, a copy of the `sentinel` value is also pushed to -/// `staticVec`. This is useful to extract mixed static and dynamic entries that -/// come from an AttrSizedOperandSegments trait. -static void dispatchIndexOpFoldResult(OpFoldResult ofr, - SmallVectorImpl &dynamicVec, - SmallVectorImpl &staticVec, - int64_t sentinel) { - if (auto v = ofr.dyn_cast()) { - dynamicVec.push_back(v); - staticVec.push_back(sentinel); - return; - } - APInt apInt = ofr.dyn_cast().cast().getValue(); - staticVec.push_back(apInt.getSExtValue()); -} - -static void dispatchIndexOpFoldResults(ArrayRef ofrs, - SmallVectorImpl &dynamicVec, - SmallVectorImpl &staticVec, - int64_t sentinel) { - for (auto ofr : ofrs) - dispatchIndexOpFoldResult(ofr, dynamicVec, staticVec, sentinel); -} - /// If ofr is a constant integer, i.e., an IntegerAttr or a ConstantOp with an /// IntegerAttr, return the integer. llvm::Optional mlir::getConstantIntValue(OpFoldResult ofr) { @@ -227,7 +200,6 @@ } void StandardOpsDialect::initialize() { - getContext()->loadDialect(); addOperations< #define GET_OP_LIST #include "mlir/Dialect/StandardOps/IR/Ops.cpp.inc" @@ -286,14 +258,6 @@ [](APInt a, APInt b) { return a + b; }); } -/// Extract int64_t values from the assumed ArrayAttr of IntegerAttr. -static SmallVector extractFromI64ArrayAttr(Attribute attr) { - return llvm::to_vector<4>( - llvm::map_range(attr.cast(), [](Attribute a) -> int64_t { - return a.cast().getInt(); - })); -} - /// Canonicalize a sum of a constant and (constant - something) to simply be /// a sum of constants minus something. This transformation does similar /// transformations for additions of a constant with a subtract/add of @@ -2082,499 +2046,6 @@ return areVectorCastSimpleCompatible(a, b, areCastCompatible); } -//===----------------------------------------------------------------------===// -// SubTensorOp -//===----------------------------------------------------------------------===// - -/// A subtensor result type can be fully inferred from the source type and the -/// static representation of offsets, sizes and strides. Special sentinels -/// encode the dynamic case. -Type SubTensorOp::inferResultType(RankedTensorType sourceRankedTensorType, - ArrayRef leadingStaticOffsets, - ArrayRef leadingStaticSizes, - ArrayRef leadingStaticStrides) { - // A subtensor may specify only a leading subset of offset/sizes/strides in - // which case we complete with offset=0, sizes from memref type and strides=1. - unsigned rank = sourceRankedTensorType.getRank(); - assert(leadingStaticSizes.size() <= rank && - "unexpected leadingStaticSizes overflow"); - auto staticSizes = llvm::to_vector<4>(leadingStaticSizes); - unsigned numTrailingSizes = rank - staticSizes.size(); - llvm::append_range(staticSizes, sourceRankedTensorType.getShape().take_back( - numTrailingSizes)); - return RankedTensorType::get(staticSizes, - sourceRankedTensorType.getElementType()); -} - -Type SubTensorOp::inferResultType(RankedTensorType sourceRankedTensorType, - ArrayRef leadingStaticOffsets, - ArrayRef leadingStaticSizes, - ArrayRef leadingStaticStrides) { - SmallVector staticOffsets, staticSizes, staticStrides; - SmallVector dynamicOffsets, dynamicSizes, dynamicStrides; - dispatchIndexOpFoldResults(leadingStaticOffsets, dynamicOffsets, - staticOffsets, ShapedType::kDynamicStrideOrOffset); - dispatchIndexOpFoldResults(leadingStaticSizes, dynamicSizes, staticSizes, - ShapedType::kDynamicSize); - dispatchIndexOpFoldResults(leadingStaticStrides, dynamicStrides, - staticStrides, ShapedType::kDynamicStrideOrOffset); - return SubTensorOp::inferResultType(sourceRankedTensorType, staticOffsets, - staticSizes, staticStrides); -} - -/// A subtensor result type can be fully inferred from the source type and the -/// static representation of offsets, sizes and strides. Special sentinels -/// encode the dynamic case. -Type SubTensorOp::inferRankReducedResultType( - unsigned resultRank, RankedTensorType sourceRankedTensorType, - ArrayRef leadingStaticOffsets, - ArrayRef leadingStaticSizes, - ArrayRef leadingStaticStrides) { - auto inferredType = - inferResultType(sourceRankedTensorType, leadingStaticOffsets, - leadingStaticSizes, leadingStaticStrides) - .cast(); - int rankDiff = inferredType.getRank() - resultRank; - if (rankDiff > 0) { - auto shape = inferredType.getShape(); - llvm::SmallDenseSet dimsToProject; - mlir::getPositionsOfShapeOne(rankDiff, shape, dimsToProject); - SmallVector projectedShape; - for (unsigned pos = 0, e = shape.size(); pos < e; ++pos) - if (!dimsToProject.contains(pos)) - projectedShape.push_back(shape[pos]); - inferredType = - RankedTensorType::get(projectedShape, inferredType.getElementType()); - } - return inferredType; -} - -Type SubTensorOp::inferRankReducedResultType( - unsigned resultRank, RankedTensorType sourceRankedTensorType, - ArrayRef leadingStaticOffsets, - ArrayRef leadingStaticSizes, - ArrayRef leadingStaticStrides) { - SmallVector staticOffsets, staticSizes, staticStrides; - SmallVector dynamicOffsets, dynamicSizes, dynamicStrides; - dispatchIndexOpFoldResults(leadingStaticOffsets, dynamicOffsets, - staticOffsets, ShapedType::kDynamicStrideOrOffset); - dispatchIndexOpFoldResults(leadingStaticSizes, dynamicSizes, staticSizes, - ShapedType::kDynamicSize); - dispatchIndexOpFoldResults(leadingStaticStrides, dynamicStrides, - staticStrides, ShapedType::kDynamicStrideOrOffset); - return SubTensorOp::inferRankReducedResultType( - resultRank, sourceRankedTensorType, staticOffsets, staticSizes, - staticStrides); -} - -// Build a SubTensorOp with mixed static and dynamic entries and custom result -// type. If the type passed is nullptr, it is inferred. -void mlir::SubTensorOp::build(OpBuilder &b, OperationState &result, - RankedTensorType resultType, Value source, - ArrayRef offsets, - ArrayRef sizes, - ArrayRef strides, - ArrayRef attrs) { - SmallVector staticOffsets, staticSizes, staticStrides; - SmallVector dynamicOffsets, dynamicSizes, dynamicStrides; - dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets, - ShapedType::kDynamicStrideOrOffset); - dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes, - ShapedType::kDynamicSize); - dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides, - ShapedType::kDynamicStrideOrOffset); - auto sourceRankedTensorType = source.getType().cast(); - // Structuring implementation this way avoids duplication between builders. - if (!resultType) { - resultType = - SubTensorOp::inferResultType(sourceRankedTensorType, staticOffsets, - staticSizes, staticStrides) - .cast(); - } - build(b, result, resultType, source, dynamicOffsets, dynamicSizes, - dynamicStrides, b.getI64ArrayAttr(staticOffsets), - b.getI64ArrayAttr(staticSizes), b.getI64ArrayAttr(staticStrides)); - result.addAttributes(attrs); -} - -// Build a SubTensorOp with mixed static and dynamic entries and inferred result -// type. -void mlir::SubTensorOp::build(OpBuilder &b, OperationState &result, - Value source, ArrayRef offsets, - ArrayRef sizes, - ArrayRef strides, - ArrayRef attrs) { - build(b, result, RankedTensorType(), source, offsets, sizes, strides, attrs); -} - -// Build a SubTensorOp with dynamic entries and custom result type. If the type -// passed is nullptr, it is inferred. -void mlir::SubTensorOp::build(OpBuilder &b, OperationState &result, - RankedTensorType resultType, Value source, - ValueRange offsets, ValueRange sizes, - ValueRange strides, - ArrayRef attrs) { - SmallVector offsetValues = llvm::to_vector<4>( - llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; })); - SmallVector sizeValues = llvm::to_vector<4>( - llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; })); - SmallVector strideValues = llvm::to_vector<4>( - llvm::map_range(strides, [](Value v) -> OpFoldResult { return v; })); - build(b, result, resultType, source, offsetValues, sizeValues, strideValues); -} - -// Build a SubTensorOp with dynamic entries and inferred result type. -void mlir::SubTensorOp::build(OpBuilder &b, OperationState &result, - Value source, ValueRange offsets, - ValueRange sizes, ValueRange strides, - ArrayRef attrs) { - build(b, result, RankedTensorType(), source, offsets, sizes, strides, attrs); -} - -enum SubTensorVerificationResult { - Success, - RankTooLarge, - SizeMismatch, - ElemTypeMismatch, -}; - -/// Checks if `original` Type type can be rank reduced to `reduced` type. -/// This function is slight variant of `is subsequence` algorithm where -/// not matching dimension must be 1. -static SubTensorVerificationResult -isRankReducedType(Type originalType, Type candidateReducedType, - std::string *errMsg = nullptr) { - if (originalType == candidateReducedType) - return SubTensorVerificationResult::Success; - if (!originalType.isa()) - return SubTensorVerificationResult::Success; - if (originalType.isa() && - !candidateReducedType.isa()) - return SubTensorVerificationResult::Success; - - ShapedType originalShapedType = originalType.cast(); - ShapedType candidateReducedShapedType = - candidateReducedType.cast(); - - // Rank and size logic is valid for all ShapedTypes. - ArrayRef originalShape = originalShapedType.getShape(); - ArrayRef candidateReducedShape = - candidateReducedShapedType.getShape(); - unsigned originalRank = originalShape.size(), - candidateReducedRank = candidateReducedShape.size(); - if (candidateReducedRank > originalRank) - return SubTensorVerificationResult::RankTooLarge; - - auto optionalUnusedDimsMask = - computeRankReductionMask(originalShape, candidateReducedShape); - - // Sizes cannot be matched in case empty vector is returned. - if (!optionalUnusedDimsMask.hasValue()) - return SubTensorVerificationResult::SizeMismatch; - - if (originalShapedType.getElementType() != - candidateReducedShapedType.getElementType()) - return SubTensorVerificationResult::ElemTypeMismatch; - - // We are done for the tensor case. - if (originalType.isa()) - return SubTensorVerificationResult::Success; - - return SubTensorVerificationResult::Success; -} - -template -static LogicalResult -produceSubTensorErrorMsg(SubTensorVerificationResult result, OpTy op, - Type expectedType, StringRef errMsg = "") { - auto memrefType = expectedType.cast(); - switch (result) { - case SubTensorVerificationResult::Success: - return success(); - case SubTensorVerificationResult::RankTooLarge: - return op.emitError("expected result rank to be smaller or equal to ") - << "the source rank. " << errMsg; - case SubTensorVerificationResult::SizeMismatch: - return op.emitError("expected result type to be ") - << expectedType - << " or a rank-reduced version. (mismatch of result sizes) " - << errMsg; - case SubTensorVerificationResult::ElemTypeMismatch: - return op.emitError("expected result element type to be ") - << memrefType.getElementType() << errMsg; - } - llvm_unreachable("unexpected subtensor verification result"); -} -/// Verifier for SubTensorOp. -static LogicalResult verify(SubTensorOp op) { - // Verify result type against inferred type. - auto expectedType = SubTensorOp::inferResultType( - op.getSourceType(), extractFromI64ArrayAttr(op.static_offsets()), - extractFromI64ArrayAttr(op.static_sizes()), - extractFromI64ArrayAttr(op.static_strides())); - auto result = isRankReducedType(expectedType, op.getType()); - return produceSubTensorErrorMsg(result, op, expectedType); -} - -/// Infer the canonical type of the result of a subtensor operation. Returns a -/// type with rank `resultRank` that is either the rank of the rank-reduced -/// type, or the non-rank-reduced type. -static RankedTensorType getCanonicalSubTensorResultType( - unsigned resultRank, RankedTensorType sourceType, - ArrayRef mixedOffsets, ArrayRef mixedSizes, - ArrayRef mixedStrides) { - auto resultType = - SubTensorOp::inferRankReducedResultType( - resultRank, sourceType, mixedOffsets, mixedSizes, mixedStrides) - .cast(); - if (resultType.getRank() != resultRank) { - resultType = SubTensorOp::inferResultType(sourceType, mixedOffsets, - mixedSizes, mixedStrides) - .cast(); - } - return resultType; -} - -namespace { -/// Pattern to rewrite a subtensor op with tensor::Cast arguments. -/// This essentially pushes memref_cast past its consuming subtensor when -/// `canFoldIntoConsumerOp` is true. -/// -/// Example: -/// ``` -/// %0 = tensorcast %V : tensor<16x16xf32> to tensor -/// %1 = subtensor %0[0, 0][3, 4][1, 1] : tensor to tensor<3x4xf32> -/// ``` -/// is rewritten into: -/// ``` -/// %0 = subtensor %V[0, 0][3, 4][1, 1] : tensor<16x16xf32> to tensor<3x4xf32> -/// %1 = tensor.cast %0: tensor<3x4xf32> to tensor<3x4xf32> -/// ``` -class SubTensorOpCastFolder final : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(SubTensorOp subTensorOp, - PatternRewriter &rewriter) const override { - // Any constant operand, just return to let SubViewOpConstantFolder kick in. - if (llvm::any_of(subTensorOp.getOperands(), [](Value operand) { - return matchPattern(operand, matchConstantIndex()); - })) - return failure(); - - auto castOp = subTensorOp.source().getDefiningOp(); - if (!castOp) - return failure(); - - if (!canFoldIntoConsumerOp(castOp)) - return failure(); - - /// Deduce the type of the result to use for the canonicalized operation. - RankedTensorType resultType = getCanonicalSubTensorResultType( - subTensorOp.getType().getRank(), subTensorOp.getSourceType(), - subTensorOp.getMixedOffsets(), subTensorOp.getMixedSizes(), - subTensorOp.getMixedStrides()); - Value newSubTensor = rewriter.create( - subTensorOp.getLoc(), resultType, castOp.source(), - subTensorOp.offsets(), subTensorOp.sizes(), subTensorOp.strides(), - subTensorOp.static_offsets(), subTensorOp.static_sizes(), - subTensorOp.static_strides()); - rewriter.replaceOpWithNewOp( - subTensorOp, subTensorOp.getType(), newSubTensor); - return success(); - } -}; -} // namespace - -/// Return the canonical type of the result of a subtensor. -struct SubTensorReturnTypeCanonicalizer { - RankedTensorType operator()(SubTensorOp op, - ArrayRef mixedOffsets, - ArrayRef mixedSizes, - ArrayRef mixedStrides) { - return getCanonicalSubTensorResultType(op.getType().getRank(), - op.getSourceType(), mixedOffsets, - mixedSizes, mixedStrides); - } -}; - -/// A canonicalizer wrapper to replace SubTensorOps. -struct SubTensorCanonicalizer { - void operator()(PatternRewriter &rewriter, SubTensorOp op, - SubTensorOp newOp) { - Value replacement = newOp.getResult(); - if (replacement.getType() != op.getType()) - replacement = rewriter.create(op.getLoc(), op.getType(), - replacement); - rewriter.replaceOp(op, replacement); - } -}; - -void SubTensorOp::getCanonicalizationPatterns(RewritePatternSet &results, - MLIRContext *context) { - results.add, - SubTensorOpCastFolder>(context); -} - -// -static LogicalResult -foldIdentityOffsetSizeAndStrideOpInterface(OffsetSizeAndStrideOpInterface op, - ShapedType shapedType) { - OpBuilder b(op.getContext()); - for (OpFoldResult ofr : op.getMixedOffsets()) - if (!isEqualConstantIntOrValue(ofr, b.getIndexAttr(0))) - return failure(); - // Rank-reducing noops only need to inspect the leading dimensions: llvm::zip - // is appropriate. - auto shape = shapedType.getShape(); - for (auto it : llvm::zip(op.getMixedSizes(), shape)) - if (!isEqualConstantIntOrValue(std::get<0>(it), - b.getIndexAttr(std::get<1>(it)))) - return failure(); - for (OpFoldResult ofr : op.getMixedStrides()) - if (!isEqualConstantIntOrValue(ofr, b.getIndexAttr(1))) - return failure(); - return success(); -} - -OpFoldResult SubTensorOp::fold(ArrayRef) { - if (getSourceType() == getType() && - succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType()))) - return this->source(); - return OpFoldResult(); -} - -//===----------------------------------------------------------------------===// -// SubTensorInsertOp -//===----------------------------------------------------------------------===// - -// Build a SubTensorInsertOp with mixed static and dynamic entries. -void mlir::SubTensorInsertOp::build(OpBuilder &b, OperationState &result, - Value source, Value dest, - ArrayRef offsets, - ArrayRef sizes, - ArrayRef strides, - ArrayRef attrs) { - SmallVector staticOffsets, staticSizes, staticStrides; - SmallVector dynamicOffsets, dynamicSizes, dynamicStrides; - dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets, - ShapedType::kDynamicStrideOrOffset); - dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes, - ShapedType::kDynamicSize); - dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides, - ShapedType::kDynamicStrideOrOffset); - build(b, result, dest.getType(), source, dest, dynamicOffsets, dynamicSizes, - dynamicStrides, b.getI64ArrayAttr(staticOffsets), - b.getI64ArrayAttr(staticSizes), b.getI64ArrayAttr(staticStrides)); - result.addAttributes(attrs); -} - -// Build a SubTensorInsertOp with dynamic entries. -void mlir::SubTensorInsertOp::build(OpBuilder &b, OperationState &result, - Value source, Value dest, - ValueRange offsets, ValueRange sizes, - ValueRange strides, - ArrayRef attrs) { - SmallVector offsetValues = llvm::to_vector<4>( - llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; })); - SmallVector sizeValues = llvm::to_vector<4>( - llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; })); - SmallVector strideValues = llvm::to_vector<4>( - llvm::map_range(strides, [](Value v) -> OpFoldResult { return v; })); - build(b, result, source, dest, offsetValues, sizeValues, strideValues); -} - -OpFoldResult SubTensorInsertOp::fold(ArrayRef) { - if (getSourceType().hasStaticShape() && getType().hasStaticShape() && - getSourceType() == getType() && - succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType()))) - return this->source(); - return OpFoldResult(); -} - -namespace { -/// Pattern to rewrite a subtensor_insert op with constant arguments. -class SubTensorInsertOpConstantArgumentFolder final - : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(SubTensorInsertOp subTensorInsertOp, - PatternRewriter &rewriter) const override { - // No constant operand, just return. - if (llvm::none_of(subTensorInsertOp.getOperands(), [](Value operand) { - return matchPattern(operand, matchConstantIndex()); - })) - return failure(); - - // At least one of offsets/sizes/strides is a new constant. - // Form the new list of operands and constant attributes from the - // existing. - SmallVector mixedOffsets(subTensorInsertOp.getMixedOffsets()); - SmallVector mixedSizes(subTensorInsertOp.getMixedSizes()); - SmallVector mixedStrides(subTensorInsertOp.getMixedStrides()); - canonicalizeSubViewPart(mixedOffsets, ShapedType::isDynamicStrideOrOffset); - canonicalizeSubViewPart(mixedSizes, ShapedType::isDynamic); - canonicalizeSubViewPart(mixedStrides, ShapedType::isDynamicStrideOrOffset); - - // Create the new op in canonical form. - rewriter.replaceOpWithNewOp( - subTensorInsertOp, subTensorInsertOp.source(), subTensorInsertOp.dest(), - mixedOffsets, mixedSizes, mixedStrides); - return success(); - } -}; - -/// Fold tensor_casts with subtensor_insert operations. -struct SubTensorInsertOpCastFolder final - : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(SubTensorInsertOp subTensorInsertOp, - PatternRewriter &rewriter) const override { - if (llvm::any_of(subTensorInsertOp.getOperands(), [](Value operand) { - return matchPattern(operand, matchConstantIndex()); - })) - return failure(); - - auto getSourceOfCastOp = [](Value v) -> Optional { - auto castOp = v.getDefiningOp(); - if (!castOp || !canFoldIntoConsumerOp(castOp)) - return llvm::None; - return castOp.source(); - }; - Optional sourceCastSource = - getSourceOfCastOp(subTensorInsertOp.source()); - Optional destCastSource = - getSourceOfCastOp(subTensorInsertOp.dest()); - if (!sourceCastSource && !destCastSource) - return failure(); - - Value replacement = rewriter.create( - subTensorInsertOp.getLoc(), - (sourceCastSource ? *sourceCastSource : subTensorInsertOp.source()), - (destCastSource ? *destCastSource : subTensorInsertOp.dest()), - subTensorInsertOp.getMixedOffsets(), subTensorInsertOp.getMixedSizes(), - subTensorInsertOp.getMixedStrides()); - - if (replacement.getType() != subTensorInsertOp.getType()) { - replacement = rewriter.create( - subTensorInsertOp.getLoc(), subTensorInsertOp.getType(), replacement); - } - rewriter.replaceOp(subTensorInsertOp, replacement); - return success(); - } -}; -} // namespace - -void SubTensorInsertOp::getCanonicalizationPatterns(RewritePatternSet &results, - MLIRContext *context) { - results.add(context); -} - //===----------------------------------------------------------------------===// // SwitchOp //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp --- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp +++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp @@ -6,6 +6,7 @@ // //===----------------------------------------------------------------------===// +#include "mlir/Dialect/StandardOps/Utils/Utils.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/BlockAndValueMapping.h" #include "mlir/IR/Builders.h" @@ -511,6 +512,530 @@ return success(); } +//===----------------------------------------------------------------------===// +// SubTensorOp +//===----------------------------------------------------------------------===// + +/// Helper function to dispatch an OpFoldResult into either the `dynamicVec` if +/// it is a Value or into `staticVec` if it is an IntegerAttr. +/// In the case of a Value, a copy of the `sentinel` value is also pushed to +/// `staticVec`. This is useful to extract mixed static and dynamic entries that +/// come from an AttrSizedOperandSegments trait. +static void dispatchIndexOpFoldResult(OpFoldResult ofr, + SmallVectorImpl &dynamicVec, + SmallVectorImpl &staticVec, + int64_t sentinel) { + if (auto v = ofr.dyn_cast()) { + dynamicVec.push_back(v); + staticVec.push_back(sentinel); + return; + } + APInt apInt = ofr.dyn_cast().cast().getValue(); + staticVec.push_back(apInt.getSExtValue()); +} + +static void dispatchIndexOpFoldResults(ArrayRef ofrs, + SmallVectorImpl &dynamicVec, + SmallVectorImpl &staticVec, + int64_t sentinel) { + for (auto ofr : ofrs) + dispatchIndexOpFoldResult(ofr, dynamicVec, staticVec, sentinel); +} + +/// A subtensor result type can be fully inferred from the source type and the +/// static representation of offsets, sizes and strides. Special sentinels +/// encode the dynamic case. +Type SubTensorOp::inferResultType(RankedTensorType sourceRankedTensorType, + ArrayRef leadingStaticOffsets, + ArrayRef leadingStaticSizes, + ArrayRef leadingStaticStrides) { + // A subtensor may specify only a leading subset of offset/sizes/strides in + // which case we complete with offset=0, sizes from memref type and strides=1. + unsigned rank = sourceRankedTensorType.getRank(); + assert(leadingStaticSizes.size() <= rank && + "unexpected leadingStaticSizes overflow"); + auto staticSizes = llvm::to_vector<4>(leadingStaticSizes); + unsigned numTrailingSizes = rank - staticSizes.size(); + llvm::append_range(staticSizes, sourceRankedTensorType.getShape().take_back( + numTrailingSizes)); + return RankedTensorType::get(staticSizes, + sourceRankedTensorType.getElementType()); +} + +/// Extract int64_t values from the assumed ArrayAttr of IntegerAttr. +static SmallVector extractFromI64ArrayAttr(Attribute attr) { + return llvm::to_vector<4>( + llvm::map_range(attr.cast(), [](Attribute a) -> int64_t { + return a.cast().getInt(); + })); +} + +Type SubTensorOp::inferResultType(RankedTensorType sourceRankedTensorType, + ArrayRef leadingStaticOffsets, + ArrayRef leadingStaticSizes, + ArrayRef leadingStaticStrides) { + SmallVector staticOffsets, staticSizes, staticStrides; + SmallVector dynamicOffsets, dynamicSizes, dynamicStrides; + dispatchIndexOpFoldResults(leadingStaticOffsets, dynamicOffsets, + staticOffsets, ShapedType::kDynamicStrideOrOffset); + dispatchIndexOpFoldResults(leadingStaticSizes, dynamicSizes, staticSizes, + ShapedType::kDynamicSize); + dispatchIndexOpFoldResults(leadingStaticStrides, dynamicStrides, + staticStrides, ShapedType::kDynamicStrideOrOffset); + return SubTensorOp::inferResultType(sourceRankedTensorType, staticOffsets, + staticSizes, staticStrides); +} + +/// A subtensor result type can be fully inferred from the source type and the +/// static representation of offsets, sizes and strides. Special sentinels +/// encode the dynamic case. +Type SubTensorOp::inferRankReducedResultType( + unsigned resultRank, RankedTensorType sourceRankedTensorType, + ArrayRef leadingStaticOffsets, + ArrayRef leadingStaticSizes, + ArrayRef leadingStaticStrides) { + auto inferredType = + inferResultType(sourceRankedTensorType, leadingStaticOffsets, + leadingStaticSizes, leadingStaticStrides) + .cast(); + int rankDiff = inferredType.getRank() - resultRank; + if (rankDiff > 0) { + auto shape = inferredType.getShape(); + llvm::SmallDenseSet dimsToProject; + mlir::getPositionsOfShapeOne(rankDiff, shape, dimsToProject); + SmallVector projectedShape; + for (unsigned pos = 0, e = shape.size(); pos < e; ++pos) + if (!dimsToProject.contains(pos)) + projectedShape.push_back(shape[pos]); + inferredType = + RankedTensorType::get(projectedShape, inferredType.getElementType()); + } + return inferredType; +} + +Type SubTensorOp::inferRankReducedResultType( + unsigned resultRank, RankedTensorType sourceRankedTensorType, + ArrayRef leadingStaticOffsets, + ArrayRef leadingStaticSizes, + ArrayRef leadingStaticStrides) { + SmallVector staticOffsets, staticSizes, staticStrides; + SmallVector dynamicOffsets, dynamicSizes, dynamicStrides; + dispatchIndexOpFoldResults(leadingStaticOffsets, dynamicOffsets, + staticOffsets, ShapedType::kDynamicStrideOrOffset); + dispatchIndexOpFoldResults(leadingStaticSizes, dynamicSizes, staticSizes, + ShapedType::kDynamicSize); + dispatchIndexOpFoldResults(leadingStaticStrides, dynamicStrides, + staticStrides, ShapedType::kDynamicStrideOrOffset); + return SubTensorOp::inferRankReducedResultType( + resultRank, sourceRankedTensorType, staticOffsets, staticSizes, + staticStrides); +} + +// Build a SubTensorOp with mixed static and dynamic entries and custom result +// type. If the type passed is nullptr, it is inferred. +void SubTensorOp::build(OpBuilder &b, OperationState &result, + RankedTensorType resultType, Value source, + ArrayRef offsets, + ArrayRef sizes, + ArrayRef strides, + ArrayRef attrs) { + SmallVector staticOffsets, staticSizes, staticStrides; + SmallVector dynamicOffsets, dynamicSizes, dynamicStrides; + dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets, + ShapedType::kDynamicStrideOrOffset); + dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes, + ShapedType::kDynamicSize); + dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides, + ShapedType::kDynamicStrideOrOffset); + auto sourceRankedTensorType = source.getType().cast(); + // Structuring implementation this way avoids duplication between builders. + if (!resultType) { + resultType = + SubTensorOp::inferResultType(sourceRankedTensorType, staticOffsets, + staticSizes, staticStrides) + .cast(); + } + build(b, result, resultType, source, dynamicOffsets, dynamicSizes, + dynamicStrides, b.getI64ArrayAttr(staticOffsets), + b.getI64ArrayAttr(staticSizes), b.getI64ArrayAttr(staticStrides)); + result.addAttributes(attrs); +} + +// Build a SubTensorOp with mixed static and dynamic entries and inferred result +// type. +void SubTensorOp::build(OpBuilder &b, OperationState &result, Value source, + ArrayRef offsets, + ArrayRef sizes, + ArrayRef strides, + ArrayRef attrs) { + build(b, result, RankedTensorType(), source, offsets, sizes, strides, attrs); +} + +// Build a SubTensorOp with dynamic entries and custom result type. If the type +// passed is nullptr, it is inferred. +void SubTensorOp::build(OpBuilder &b, OperationState &result, + RankedTensorType resultType, Value source, + ValueRange offsets, ValueRange sizes, + ValueRange strides, ArrayRef attrs) { + SmallVector offsetValues = llvm::to_vector<4>( + llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; })); + SmallVector sizeValues = llvm::to_vector<4>( + llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; })); + SmallVector strideValues = llvm::to_vector<4>( + llvm::map_range(strides, [](Value v) -> OpFoldResult { return v; })); + build(b, result, resultType, source, offsetValues, sizeValues, strideValues); +} + +// Build a SubTensorOp with dynamic entries and inferred result type. +void SubTensorOp::build(OpBuilder &b, OperationState &result, Value source, + ValueRange offsets, ValueRange sizes, + ValueRange strides, ArrayRef attrs) { + build(b, result, RankedTensorType(), source, offsets, sizes, strides, attrs); +} + +enum SubTensorVerificationResult { + Success, + RankTooLarge, + SizeMismatch, + ElemTypeMismatch, +}; + +/// Checks if `original` Type type can be rank reduced to `reduced` type. +/// This function is slight variant of `is subsequence` algorithm where +/// not matching dimension must be 1. +static SubTensorVerificationResult +isRankReducedType(Type originalType, Type candidateReducedType, + std::string *errMsg = nullptr) { + if (originalType == candidateReducedType) + return SubTensorVerificationResult::Success; + if (!originalType.isa()) + return SubTensorVerificationResult::Success; + if (originalType.isa() && + !candidateReducedType.isa()) + return SubTensorVerificationResult::Success; + + ShapedType originalShapedType = originalType.cast(); + ShapedType candidateReducedShapedType = + candidateReducedType.cast(); + + // Rank and size logic is valid for all ShapedTypes. + ArrayRef originalShape = originalShapedType.getShape(); + ArrayRef candidateReducedShape = + candidateReducedShapedType.getShape(); + unsigned originalRank = originalShape.size(), + candidateReducedRank = candidateReducedShape.size(); + if (candidateReducedRank > originalRank) + return SubTensorVerificationResult::RankTooLarge; + + auto optionalUnusedDimsMask = + computeRankReductionMask(originalShape, candidateReducedShape); + + // Sizes cannot be matched in case empty vector is returned. + if (!optionalUnusedDimsMask.hasValue()) + return SubTensorVerificationResult::SizeMismatch; + + if (originalShapedType.getElementType() != + candidateReducedShapedType.getElementType()) + return SubTensorVerificationResult::ElemTypeMismatch; + + // We are done for the tensor case. + if (originalType.isa()) + return SubTensorVerificationResult::Success; + + return SubTensorVerificationResult::Success; +} + +template +static LogicalResult +produceSubTensorErrorMsg(SubTensorVerificationResult result, OpTy op, + Type expectedType, StringRef errMsg = "") { + auto memrefType = expectedType.cast(); + switch (result) { + case SubTensorVerificationResult::Success: + return success(); + case SubTensorVerificationResult::RankTooLarge: + return op.emitError("expected result rank to be smaller or equal to ") + << "the source rank. " << errMsg; + case SubTensorVerificationResult::SizeMismatch: + return op.emitError("expected result type to be ") + << expectedType + << " or a rank-reduced version. (mismatch of result sizes) " + << errMsg; + case SubTensorVerificationResult::ElemTypeMismatch: + return op.emitError("expected result element type to be ") + << memrefType.getElementType() << errMsg; + } + llvm_unreachable("unexpected subtensor verification result"); +} +/// Verifier for SubTensorOp. +static LogicalResult verify(SubTensorOp op) { + // Verify result type against inferred type. + auto expectedType = SubTensorOp::inferResultType( + op.getSourceType(), extractFromI64ArrayAttr(op.static_offsets()), + extractFromI64ArrayAttr(op.static_sizes()), + extractFromI64ArrayAttr(op.static_strides())); + auto result = isRankReducedType(expectedType, op.getType()); + return produceSubTensorErrorMsg(result, op, expectedType); +} + +/// Infer the canonical type of the result of a subtensor operation. Returns a +/// type with rank `resultRank` that is either the rank of the rank-reduced +/// type, or the non-rank-reduced type. +static RankedTensorType getCanonicalSubTensorResultType( + unsigned resultRank, RankedTensorType sourceType, + ArrayRef mixedOffsets, ArrayRef mixedSizes, + ArrayRef mixedStrides) { + auto resultType = + SubTensorOp::inferRankReducedResultType( + resultRank, sourceType, mixedOffsets, mixedSizes, mixedStrides) + .cast(); + if (resultType.getRank() != resultRank) { + resultType = SubTensorOp::inferResultType(sourceType, mixedOffsets, + mixedSizes, mixedStrides) + .cast(); + } + return resultType; +} + +namespace { +/// Pattern to rewrite a subtensor op with tensor::Cast arguments. +/// This essentially pushes memref_cast past its consuming subtensor when +/// `canFoldIntoConsumerOp` is true. +/// +/// Example: +/// ``` +/// %0 = tensorcast %V : tensor<16x16xf32> to tensor +/// %1 = subtensor %0[0, 0][3, 4][1, 1] : tensor to tensor<3x4xf32> +/// ``` +/// is rewritten into: +/// ``` +/// %0 = subtensor %V[0, 0][3, 4][1, 1] : tensor<16x16xf32> to tensor<3x4xf32> +/// %1 = tensor.cast %0: tensor<3x4xf32> to tensor<3x4xf32> +/// ``` +class SubTensorOpCastFolder final : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(SubTensorOp subTensorOp, + PatternRewriter &rewriter) const override { + // Any constant operand, just return to let SubViewOpConstantFolder kick in. + if (llvm::any_of(subTensorOp.getOperands(), [](Value operand) { + return matchPattern(operand, matchConstantIndex()); + })) + return failure(); + + auto castOp = subTensorOp.source().getDefiningOp(); + if (!castOp) + return failure(); + + if (!canFoldIntoConsumerOp(castOp)) + return failure(); + + /// Deduce the type of the result to use for the canonicalized operation. + RankedTensorType resultType = getCanonicalSubTensorResultType( + subTensorOp.getType().getRank(), subTensorOp.getSourceType(), + subTensorOp.getMixedOffsets(), subTensorOp.getMixedSizes(), + subTensorOp.getMixedStrides()); + Value newSubTensor = rewriter.create( + subTensorOp.getLoc(), resultType, castOp.source(), + subTensorOp.offsets(), subTensorOp.sizes(), subTensorOp.strides(), + subTensorOp.static_offsets(), subTensorOp.static_sizes(), + subTensorOp.static_strides()); + rewriter.replaceOpWithNewOp( + subTensorOp, subTensorOp.getType(), newSubTensor); + return success(); + } +}; +} // namespace + +/// Return the canonical type of the result of a subtensor. +struct SubTensorReturnTypeCanonicalizer { + RankedTensorType operator()(SubTensorOp op, + ArrayRef mixedOffsets, + ArrayRef mixedSizes, + ArrayRef mixedStrides) { + return getCanonicalSubTensorResultType(op.getType().getRank(), + op.getSourceType(), mixedOffsets, + mixedSizes, mixedStrides); + } +}; + +/// A canonicalizer wrapper to replace SubTensorOps. +struct SubTensorCanonicalizer { + void operator()(PatternRewriter &rewriter, SubTensorOp op, + SubTensorOp newOp) { + Value replacement = newOp.getResult(); + if (replacement.getType() != op.getType()) + replacement = rewriter.create(op.getLoc(), op.getType(), + replacement); + rewriter.replaceOp(op, replacement); + } +}; + +void SubTensorOp::getCanonicalizationPatterns(RewritePatternSet &results, + MLIRContext *context) { + results.add, + SubTensorOpCastFolder>(context); +} + +// +static LogicalResult +foldIdentityOffsetSizeAndStrideOpInterface(OffsetSizeAndStrideOpInterface op, + ShapedType shapedType) { + OpBuilder b(op.getContext()); + for (OpFoldResult ofr : op.getMixedOffsets()) + if (!isEqualConstantIntOrValue(ofr, b.getIndexAttr(0))) + return failure(); + // Rank-reducing noops only need to inspect the leading dimensions: llvm::zip + // is appropriate. + auto shape = shapedType.getShape(); + for (auto it : llvm::zip(op.getMixedSizes(), shape)) + if (!isEqualConstantIntOrValue(std::get<0>(it), + b.getIndexAttr(std::get<1>(it)))) + return failure(); + for (OpFoldResult ofr : op.getMixedStrides()) + if (!isEqualConstantIntOrValue(ofr, b.getIndexAttr(1))) + return failure(); + return success(); +} + +OpFoldResult SubTensorOp::fold(ArrayRef) { + if (getSourceType() == getType() && + succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType()))) + return this->source(); + return OpFoldResult(); +} + +//===----------------------------------------------------------------------===// +// SubTensorInsertOp +//===----------------------------------------------------------------------===// + +// Build a SubTensorInsertOp with mixed static and dynamic entries. +void SubTensorInsertOp::build(OpBuilder &b, OperationState &result, + Value source, Value dest, + ArrayRef offsets, + ArrayRef sizes, + ArrayRef strides, + ArrayRef attrs) { + SmallVector staticOffsets, staticSizes, staticStrides; + SmallVector dynamicOffsets, dynamicSizes, dynamicStrides; + dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets, + ShapedType::kDynamicStrideOrOffset); + dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes, + ShapedType::kDynamicSize); + dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides, + ShapedType::kDynamicStrideOrOffset); + build(b, result, dest.getType(), source, dest, dynamicOffsets, dynamicSizes, + dynamicStrides, b.getI64ArrayAttr(staticOffsets), + b.getI64ArrayAttr(staticSizes), b.getI64ArrayAttr(staticStrides)); + result.addAttributes(attrs); +} + +// Build a SubTensorInsertOp with dynamic entries. +void SubTensorInsertOp::build(OpBuilder &b, OperationState &result, + Value source, Value dest, ValueRange offsets, + ValueRange sizes, ValueRange strides, + ArrayRef attrs) { + SmallVector offsetValues = llvm::to_vector<4>( + llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; })); + SmallVector sizeValues = llvm::to_vector<4>( + llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; })); + SmallVector strideValues = llvm::to_vector<4>( + llvm::map_range(strides, [](Value v) -> OpFoldResult { return v; })); + build(b, result, source, dest, offsetValues, sizeValues, strideValues); +} + +OpFoldResult SubTensorInsertOp::fold(ArrayRef) { + if (getSourceType().hasStaticShape() && getType().hasStaticShape() && + getSourceType() == getType() && + succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType()))) + return this->source(); + return OpFoldResult(); +} + +namespace { +/// Pattern to rewrite a subtensor_insert op with constant arguments. +class SubTensorInsertOpConstantArgumentFolder final + : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(SubTensorInsertOp subTensorInsertOp, + PatternRewriter &rewriter) const override { + // No constant operand, just return. + if (llvm::none_of(subTensorInsertOp.getOperands(), [](Value operand) { + return matchPattern(operand, matchConstantIndex()); + })) + return failure(); + + // At least one of offsets/sizes/strides is a new constant. + // Form the new list of operands and constant attributes from the + // existing. + SmallVector mixedOffsets(subTensorInsertOp.getMixedOffsets()); + SmallVector mixedSizes(subTensorInsertOp.getMixedSizes()); + SmallVector mixedStrides(subTensorInsertOp.getMixedStrides()); + canonicalizeSubViewPart(mixedOffsets, ShapedType::isDynamicStrideOrOffset); + canonicalizeSubViewPart(mixedSizes, ShapedType::isDynamic); + canonicalizeSubViewPart(mixedStrides, ShapedType::isDynamicStrideOrOffset); + + // Create the new op in canonical form. + rewriter.replaceOpWithNewOp( + subTensorInsertOp, subTensorInsertOp.source(), subTensorInsertOp.dest(), + mixedOffsets, mixedSizes, mixedStrides); + return success(); + } +}; + +/// Fold tensor_casts with subtensor_insert operations. +struct SubTensorInsertOpCastFolder final + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(SubTensorInsertOp subTensorInsertOp, + PatternRewriter &rewriter) const override { + if (llvm::any_of(subTensorInsertOp.getOperands(), [](Value operand) { + return matchPattern(operand, matchConstantIndex()); + })) + return failure(); + + auto getSourceOfCastOp = [](Value v) -> Optional { + auto castOp = v.getDefiningOp(); + if (!castOp || !canFoldIntoConsumerOp(castOp)) + return llvm::None; + return castOp.source(); + }; + Optional sourceCastSource = + getSourceOfCastOp(subTensorInsertOp.source()); + Optional destCastSource = + getSourceOfCastOp(subTensorInsertOp.dest()); + if (!sourceCastSource && !destCastSource) + return failure(); + + Value replacement = rewriter.create( + subTensorInsertOp.getLoc(), + (sourceCastSource ? *sourceCastSource : subTensorInsertOp.source()), + (destCastSource ? *destCastSource : subTensorInsertOp.dest()), + subTensorInsertOp.getMixedOffsets(), subTensorInsertOp.getMixedSizes(), + subTensorInsertOp.getMixedStrides()); + + if (replacement.getType() != subTensorInsertOp.getType()) { + replacement = rewriter.create( + subTensorInsertOp.getLoc(), subTensorInsertOp.getType(), replacement); + } + rewriter.replaceOp(subTensorInsertOp, replacement); + return success(); + } +}; +} // namespace + +void SubTensorInsertOp::getCanonicalizationPatterns(RewritePatternSet &results, + MLIRContext *context) { + results.add(context); +} + //===----------------------------------------------------------------------===// // TableGen'd op method definitions //===----------------------------------------------------------------------===// 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 @@ -616,9 +616,9 @@ // CHECK-NEXT: %[[ISNEG:.*]] = cmpi slt, %[[INDEX]], %[[C0]] : index // CHECK-NEXT: %[[SELECT:.*]] = select %[[ISNEG]], %[[POSINDEX]], %[[INDEX]] : index // CHECK-NEXT: %[[C1:.*]] = constant 1 : index - // CHECK-NEXT: %[[HEAD:.*]] = subtensor %[[SHAPE]][%[[C0]]] [%[[SELECT]]] [%[[C1]]] : tensor to tensor + // CHECK-NEXT: %[[HEAD:.*]] = tensor.subtensor %[[SHAPE]][%[[C0]]] [%[[SELECT]]] [%[[C1]]] : tensor to tensor // CHECK-NEXT: %[[TAIL_SIZE:.*]] = subi %[[RANK]], %[[SELECT]] : index - // CHECK-NEXT: %[[TAIL:.*]] = subtensor %[[SHAPE]][%[[SELECT]]] [%[[TAIL_SIZE]]] [%[[C1]]] : tensor to tensor + // CHECK-NEXT: %[[TAIL:.*]] = tensor.subtensor %[[SHAPE]][%[[SELECT]]] [%[[TAIL_SIZE]]] [%[[C1]]] : tensor to tensor // CHECK-NEXT: return %[[HEAD]], %[[TAIL]] : tensor, tensor %head, %tail = "shape.split_at"(%shape, %index) : (tensor, index) -> (tensor, tensor) return %head, %tail : tensor, tensor diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir --- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir +++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir @@ -679,10 +679,10 @@ // CHECK: [[CST:%.+]] = constant 0.0 // CHECK: [[FILL:%.+]] = linalg.fill([[INIT]], [[CST]]) // CHECK: [[ARG0_DIM0:%.+]] = memref.dim %arg0, [[AXIS]] - // CHECK: [[INSERT0:%.+]] = subtensor_insert %arg0 into [[FILL]]{{\[}}[[OFFSET]], [[OFFSET]]] {{\[}}[[ARG0_DIM0]], [[ARG0_DIM1]]] {{\[}}[[STRIDE]], [[STRIDE]]] + // CHECK: [[INSERT0:%.+]] = tensor.subtensor_insert %arg0 into [[FILL]]{{\[}}[[OFFSET]], [[OFFSET]]] {{\[}}[[ARG0_DIM0]], [[ARG0_DIM1]]] {{\[}}[[STRIDE]], [[STRIDE]]] // CHECK: [[NEW_OFFSET:%.+]] = addi [[OFFSET]], [[ARG0_DIM0]] // CHECK: [[ARG1_DIM0:%.+]] = memref.dim %arg1, [[AXIS]] - // CHECK: [[INSERT1:%.+]] = subtensor_insert %arg1 into [[INSERT0]]{{\[}}[[NEW_OFFSET]], [[OFFSET]]] {{\[}}[[ARG1_DIM0]], [[ARG0_DIM1]]] {{\[}}[[STRIDE]], [[STRIDE]]] + // CHECK: [[INSERT1:%.+]] = tensor.subtensor_insert %arg1 into [[INSERT0]]{{\[}}[[NEW_OFFSET]], [[OFFSET]]] {{\[}}[[ARG1_DIM0]], [[ARG0_DIM1]]] {{\[}}[[STRIDE]], [[STRIDE]]] %0 = "tosa.concat"(%arg0, %arg1) { axis = 0 : i64} : (tensor<5x1xf32>, tensor<6x1xf32>) -> (tensor<11x1xf32>) // CHECK: [[AXIS:%.+]] = constant 1 @@ -698,10 +698,10 @@ // CHECK: [[CST:%.+]] = constant 0.0 // CHECK: [[FILL:%.+]] = linalg.fill([[INIT]], [[CST]]) // CHECK: [[ARG0_DIM1:%.+]] = memref.dim %arg0, [[AXIS]] - // CHECK: [[INSERT0:%.+]] = subtensor_insert %arg0 into [[FILL]]{{\[}}[[OFFSET]], [[OFFSET]]] {{\[}}[[ARG0_DIM0]], [[ARG0_DIM1]]] {{\[}}[[STRIDE]], [[STRIDE]]] + // CHECK: [[INSERT0:%.+]] = tensor.subtensor_insert %arg0 into [[FILL]]{{\[}}[[OFFSET]], [[OFFSET]]] {{\[}}[[ARG0_DIM0]], [[ARG0_DIM1]]] {{\[}}[[STRIDE]], [[STRIDE]]] // CHECK: [[NEW_OFFSET:%.+]] = addi [[OFFSET]], [[ARG0_DIM1]] // CHECK: [[ARG1_DIM1:%.+]] = memref.dim %arg0, [[AXIS]] - // CHECK: [[INSERT1:%.+]] = subtensor_insert %arg0 into [[INSERT0]]{{\[}}[[OFFSET]], [[NEW_OFFSET]]] {{\[}}[[ARG0_DIM0]], [[ARG1_DIM1]]] {{\[}}[[STRIDE]], [[STRIDE]]] + // CHECK: [[INSERT1:%.+]] = tensor.subtensor_insert %arg0 into [[INSERT0]]{{\[}}[[OFFSET]], [[NEW_OFFSET]]] {{\[}}[[ARG0_DIM0]], [[ARG1_DIM1]]] {{\[}}[[STRIDE]], [[STRIDE]]] %1 = "tosa.concat"(%arg0, %arg0) { axis = 1 : i64} : (tensor<5x1xf32>, tensor<5x1xf32>) -> (tensor<5x2xf32>) return } diff --git a/mlir/test/Conversion/TosaToStandard/tosa-to-standard.mlir b/mlir/test/Conversion/TosaToStandard/tosa-to-standard.mlir --- a/mlir/test/Conversion/TosaToStandard/tosa-to-standard.mlir +++ b/mlir/test/Conversion/TosaToStandard/tosa-to-standard.mlir @@ -12,7 +12,7 @@ // ----- func @slice(%arg0: tensor<6xf32>) ->() { - // CHECK: [[SLICE:%.+]] = subtensor %arg0[2] [1] [1] + // CHECK: [[SLICE:%.+]] = tensor.subtensor %arg0[2] [1] [1] %0 = "tosa.slice"(%arg0) {start = [2], size = [1]} : (tensor<6xf32>) -> (tensor<1xf32>) return } diff --git a/mlir/test/Dialect/Linalg/bufferize.mlir b/mlir/test/Dialect/Linalg/bufferize.mlir --- a/mlir/test/Dialect/Linalg/bufferize.mlir +++ b/mlir/test/Dialect/Linalg/bufferize.mlir @@ -178,14 +178,14 @@ // CHECK-SAME: memref to memref<2x3xf32, #[[$MAP0]]> // CHECK-NEXT: linalg.copy(%[[SM0]], %[[A0]]) : memref<2x3xf32, #[[$MAP0]]>, memref<2x3xf32> // CHECK-NEXT: %[[RT0:.*]] = memref.tensor_load %[[A0]] : memref<2x3xf32> - %st0 = subtensor %t[0, 0][2, 3][1, 1] : tensor to tensor<2x3xf32> + %st0 = tensor.subtensor %t[0, 0][2, 3][1, 1] : tensor to tensor<2x3xf32> // CHECK-NEXT: %[[A1:.*]] = memref.alloc(%[[IDX]]) : memref<2x?xf32> // CHECK-NEXT: %[[SM1:.*]] = memref.subview %[[M]][0, %[[IDX]]] [2, %[[IDX]]] [1, 2] // CHECK-SAME: memref to memref<2x?xf32, #[[$MAP1]]> // CHECK-NEXT: linalg.copy(%[[SM1]], %[[A1]]) : memref<2x?xf32, #[[$MAP1]]>, memref<2x?xf32> // CHECK-NEXT: %[[RT1:.*]] = memref.tensor_load %[[A1]] : memref<2x?xf32> - %st1 = subtensor %t[0, %i0][2, %i0][1, 2] : tensor to tensor<2x?xf32> + %st1 = tensor.subtensor %t[0, %i0][2, %i0][1, 2] : tensor to tensor<2x?xf32> // CHECK-NEXT: return %[[RT0]], %[[RT1]] return %st0, %st1 : tensor<2x3xf32>, tensor<2x?xf32> @@ -222,7 +222,7 @@ // CHECK-SAME: memref to memref<2x3xf32, #[[$MAP0]]> // CHECK-NEXT: linalg.copy(%[[SM0]], %[[SUBVIEW0]]) : memref<2x3xf32>, memref<2x3xf32, #[[$MAP0]]> // CHECK-NEXT: %[[RT0:.*]] = memref.tensor_load %[[M_COPY0]] : memref - %t0 = subtensor_insert %st0 into %t[0, 0][2, 3][1, 1] : tensor<2x3xf32> into tensor + %t0 = tensor.subtensor_insert %st0 into %t[0, 0][2, 3][1, 1] : tensor<2x3xf32> into tensor // CHECK-DAG: %[[SM1:.*]] = memref.buffer_cast %[[ST1]] : memref<2x?xf32> // CHECK-NEXT: %[[M_COPY1:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) : memref @@ -231,7 +231,7 @@ // CHECK-SAME: memref to memref<2x?xf32, #[[$MAP1]]> // CHECK-NEXT: linalg.copy(%[[SM1]], %[[SUBVIEW1]]) : memref<2x?xf32>, memref<2x?xf32, #[[$MAP1]]> // CHECK-NEXT: %[[RT1:.*]] = memref.tensor_load %[[M_COPY1]] : memref - %t1 = subtensor_insert %st1 into %t[0, %i0][2, %i0][1, 2] : tensor<2x?xf32> into tensor + %t1 = tensor.subtensor_insert %st1 into %t[0, %i0][2, %i0][1, 2] : tensor<2x?xf32> into tensor // CHECK: return %[[RT0]], %[[RT1]] return %t0, %t1: tensor, tensor diff --git a/mlir/test/Dialect/Linalg/canonicalize.mlir b/mlir/test/Dialect/Linalg/canonicalize.mlir --- a/mlir/test/Dialect/Linalg/canonicalize.mlir +++ b/mlir/test/Dialect/Linalg/canonicalize.mlir @@ -652,7 +652,7 @@ (%arg0 : index, %arg1 : index) -> tensor<5x?x20xf32> { %0 = linalg.init_tensor[%arg0, 10, 40] : tensor - %1 = subtensor %0[0, 0, 0] [5, %arg1, 20] [1, 1, 1] + %1 = tensor.subtensor %0[0, 0, 0] [5, %arg1, 20] [1, 1, 1] : tensor to tensor<5x?x20xf32> return %1 : tensor<5x?x20xf32> } @@ -723,13 +723,13 @@ %1 = linalg.fill(%0, %arg1) : tensor, f32 -> tensor %2 = memref.dim %arg0, %c0 : tensor %3 = memref.dim %arg0, %c1 : tensor - %4 = subtensor_insert %arg0 into %1[%arg2, %arg3] [%2, %3] [1, 1] : tensor into tensor + %4 = tensor.subtensor_insert %arg0 into %1[%arg2, %arg3] [%2, %3] [1, 1] : tensor into tensor return %4 : tensor } // CHECK-LABEL: func @propogate_casts // CHECK: %[[INIT:.+]] = linalg.init_tensor [21, 42] // CHECK: %[[FILL:.+]] = linalg.fill(%[[INIT]], %{{.+}}) -// CHECK: %[[INSERTED:.+]] = subtensor_insert %{{.+}} into %[[FILL]] +// CHECK: %[[INSERTED:.+]] = tensor.subtensor_insert %{{.+}} into %[[FILL]] // CHECK: %[[RESULT:.+]] = tensor.cast %[[INSERTED]] // CHECK: return %[[RESULT]] diff --git a/mlir/test/Dialect/Linalg/comprehensive-func-bufferize.mlir b/mlir/test/Dialect/Linalg/comprehensive-func-bufferize.mlir --- a/mlir/test/Dialect/Linalg/comprehensive-func-bufferize.mlir +++ b/mlir/test/Dialect/Linalg/comprehensive-func-bufferize.mlir @@ -129,7 +129,7 @@ // CHECK-NOT: alloc // CHECK: %[[SV:.*]] = memref.subview %[[BUFFER_CAST_A]] // CHECK: linalg.copy(%[[BUFFER_CAST_B]], %[[SV]]) - %r0 = subtensor_insert %t into %A[0][4][1] : tensor<4xf32> into tensor + %r0 = tensor.subtensor_insert %t into %A[0][4][1] : tensor<4xf32> into tensor return %r0: tensor } @@ -147,7 +147,7 @@ // CHECK-NOT: alloc // CHECK: %[[SV:.*]] = memref.subview %[[BUFFER_CAST_A]] // CHECK: linalg.copy(%[[BUFFER_CAST_B]], %[[SV]]) - %r0 = subtensor_insert %t into %A[0][4][1] : tensor<4xf32> into tensor + %r0 = tensor.subtensor_insert %t into %A[0][4][1] : tensor<4xf32> into tensor /// Overwrite BUFFER_CAST_A inplace. // CHECK: linalg.fill(%[[BUFFER_CAST_A]] @@ -173,7 +173,7 @@ // CHECK: %[[SV:.*]] = memref.subview %[[BUFFER_CAST_A]] /// Overwrite BUFFER_CAST_A inplace by copying into the subview. // CHECK: linalg.copy(%[[BUFFER_CAST_B]], %[[SV]]) - %r1 = subtensor_insert %t into %r0[0][4][1] : tensor<4xf32> into tensor + %r1 = tensor.subtensor_insert %t into %r0[0][4][1] : tensor<4xf32> into tensor return %r1: tensor } @@ -192,7 +192,7 @@ // CHECK: %[[SV:.*]] = memref.subview %[[ALLOC]][0] [4] [1] : memref to memref<4xf32> // CHECK: linalg.copy(%[[BUFFER_CAST_B]], %[[SV]]) : memref<4xf32, #map>, memref<4xf32> // CHECK: memref.dealloc %[[ALLOC]] : memref - %r0 = subtensor_insert %t into %A[0][4][1] : tensor<4xf32> into tensor + %r0 = tensor.subtensor_insert %t into %A[0][4][1] : tensor<4xf32> into tensor return %r0: tensor } @@ -211,7 +211,7 @@ // CHECK: linalg.copy(%[[BUFFER_CAST_A]], %[[ALLOC]]) : memref // CHECK: %[[SV:.*]] = memref.subview %[[ALLOC]][0] [4] [1] : memref to memref<4xf32> // CHECK: linalg.copy(%[[BUFFER_CAST_B]], %[[SV]]) : memref<4xf32, #map>, memref<4xf32> - %r0 = subtensor_insert %t into %A[0][4][1] : tensor<4xf32> into tensor + %r0 = tensor.subtensor_insert %t into %A[0][4][1] : tensor<4xf32> into tensor // TODO: WAW optimization where result is overwritten without being read. // CHECK: linalg.fill(%[[BUFFER_CAST_A]] @@ -231,7 +231,7 @@ // CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<4xf32> // CHECK: %[[SV:.*]] = memref.subview %[[BUFFER_CAST_A]][0] [4] [1] // CHECK: linalg.copy(%[[SV]], %[[ALLOC]]) - %r0 = subtensor %A[0][4][1] : tensor to tensor<4xf32> + %r0 = tensor.subtensor %A[0][4][1] : tensor to tensor<4xf32> return %r0: tensor<4xf32> } @@ -243,8 +243,8 @@ %B : tensor<4x4xf32>, %C : tensor<4x4xf32>) -> tensor<4x4xf32> { // subtensor is only used as a read. - // ANALYSIS: subtensor {{.*}} {__inplace_results_attr__ = ["true"]} - %sA = subtensor %A[0, 0][4, 4][1, 1] : tensor to tensor<4x4xf32> + // ANALYSIS: tensor.subtensor {{.*}} {__inplace_results_attr__ = ["true"]} + %sA = tensor.subtensor %A[0, 0][4, 4][1, 1] : tensor to tensor<4x4xf32> // matmul output operand is not inplaceable at the function boundary. // ANALYSIS: linalg.matmul {{.*}} // ANALYSIS-NOT: {__inplace_results_attr__ = ["true"]} @@ -263,12 +263,12 @@ { // subtensor has no matching subtensor_insert and is not just used by known // readonly ops. - // ANALYSIS: subtensor {{.*}} + // ANALYSIS: tensor.subtensor {{.*}} // ANALYSIS-NOT: {__inplace_results_attr__ = ["true"]} - %r0 = subtensor %A[0][4][1] : tensor to tensor<4xf32> + %r0 = tensor.subtensor %A[0][4][1] : tensor to tensor<4xf32> // subtensor_insert can bufferize inplace fine. - // ANALYSIS: subtensor_insert {{.*}} {__inplace_results_attr__ = ["true"]} - %r1 = subtensor_insert %r0 into %A[%idx][4][1] : tensor<4xf32> into tensor + // ANALYSIS: tensor.subtensor_insert {{.*}} {__inplace_results_attr__ = ["true"]} + %r1 = tensor.subtensor_insert %r0 into %A[%idx][4][1] : tensor<4xf32> into tensor return %r1: tensor } @@ -281,12 +281,12 @@ { // subtensor has no matching subtensor_insert and is not just used by known // readonly ops. - // ANALYSIS: subtensor {{.*}} {__inplace_results_attr__ = ["true"]} - %r0 = subtensor %A[0][4][1] : tensor to tensor<4xf32> + // ANALYSIS: tensor.subtensor {{.*}} {__inplace_results_attr__ = ["true"]} + %r0 = tensor.subtensor %A[0][4][1] : tensor to tensor<4xf32> // subtensor_insert cannot bufferize inplace. - // ANALYSIS: subtensor_insert {{.*}} + // ANALYSIS: tensor.subtensor_insert {{.*}} // ANALYSIS-NOT: {__inplace_results_attr__ = ["true"]} - %r1 = subtensor_insert %r0 into %A[%idx][4][1] : tensor<4xf32> into tensor + %r1 = tensor.subtensor_insert %r0 into %A[%idx][4][1] : tensor<4xf32> into tensor return %r1: tensor } @@ -299,12 +299,12 @@ // subtensor has a matching subtensor_insert that bufferizes inplace. // TODO: Atm subtensor is not inplaceable but can be. // In the grander scheme, this will canonicalize away beforehand. - // ANALYSIS: subtensor {{.*}} + // ANALYSIS: tensor.subtensor {{.*}} // ANALYSIS-NOT: {__inplace_results_attr__ = ["true"]} - %r0 = subtensor %A[0][4][1] : tensor to tensor<4xf32> + %r0 = tensor.subtensor %A[0][4][1] : tensor to tensor<4xf32> // subtensor_insert can bufferize inplace fine. - // ANALYSIS: subtensor_insert {{.*}} {__inplace_results_attr__ = ["true"]} - %r1 = subtensor_insert %r0 into %A[0][4][1] : tensor<4xf32> into tensor + // ANALYSIS: tensor.subtensor_insert {{.*}} {__inplace_results_attr__ = ["true"]} + %r1 = tensor.subtensor_insert %r0 into %A[0][4][1] : tensor<4xf32> into tensor return %r1: tensor } @@ -318,14 +318,14 @@ // be inplaceable. // In the grander scheme, %r2 will canonicalize away beforehand but %r0 will still // not be inplaceable as the production of %r1 may involve a self-copy. - // ANALYSIS: subtensor {{.*}} + // ANALYSIS: tensor.subtensor {{.*}} // ANALYSIS-NOT: {__inplace_results_attr__ = ["true"]} - %r0 = subtensor %A[0][4][1] : tensor to tensor<4xf32> - // ANALYSIS: subtensor_insert {{.*}} + %r0 = tensor.subtensor %A[0][4][1] : tensor to tensor<4xf32> + // ANALYSIS: tensor.subtensor_insert {{.*}} // ANALYSIS-NOT: {__inplace_results_attr__ = ["true"]} - %r1 = subtensor_insert %r0 into %A[%idx][4][1] : tensor<4xf32> into tensor - // ANALYSIS: subtensor_insert {{.*}} {__inplace_results_attr__ = ["true"]} - %r2 = subtensor_insert %r0 into %A[0][4][1] : tensor<4xf32> into tensor + %r1 = tensor.subtensor_insert %r0 into %A[%idx][4][1] : tensor<4xf32> into tensor + // ANALYSIS: tensor.subtensor_insert {{.*}} {__inplace_results_attr__ = ["true"]} + %r2 = tensor.subtensor_insert %r0 into %A[0][4][1] : tensor<4xf32> into tensor return %r1, %r2: tensor, tensor } @@ -340,14 +340,14 @@ // In the grander scheme, %r2 will canonicalize away beforehand and %r0 will become // inplaceable by reducing to the `subtensor_nonmatching_subtensor_insert_non_inplace` // case, - // ANALYSIS: subtensor {{.*}} + // ANALYSIS: tensor.subtensor {{.*}} // ANALYSIS-NOT: {__inplace_results_attr__ = ["true"]} - %r0 = subtensor %A[0][4][1] : tensor to tensor<4xf32> - // ANALYSIS: subtensor_insert {{.*}} + %r0 = tensor.subtensor %A[0][4][1] : tensor to tensor<4xf32> + // ANALYSIS: tensor.subtensor_insert {{.*}} // ANALYSIS-NOT: {__inplace_results_attr__ = ["true"]} - %r2 = subtensor_insert %r0 into %A[0][4][1] : tensor<4xf32> into tensor - // ANALYSIS: subtensor_insert {{.*}} {__inplace_results_attr__ = ["true"]} - %r1 = subtensor_insert %r0 into %A[%idx][4][1] : tensor<4xf32> into tensor + %r2 = tensor.subtensor_insert %r0 into %A[0][4][1] : tensor<4xf32> into tensor + // ANALYSIS: tensor.subtensor_insert {{.*}} {__inplace_results_attr__ = ["true"]} + %r1 = tensor.subtensor_insert %r0 into %A[%idx][4][1] : tensor<4xf32> into tensor return %r1, %r2: tensor, tensor } diff --git a/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir b/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir --- a/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir +++ b/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir @@ -303,22 +303,22 @@ %arg0 : tensor<1x?x?x1x?x1x1xf32>, %arg1 : tensor<1x?x?x?x?x1x1xf32>, %arg2 : index, %arg3 : index, %arg4 : index, %arg5 : index, %arg6 : index, %arg7 : index) -> (tensor<1x?x?x1x?x1x1xf32>, tensor<1x?x?x1x?x1x1xf32>) { - %0 = subtensor %arg0[0, %arg2, %arg3, 0, %arg4, 0, 0] - [1, %arg5, %arg6, 1, %arg7, 1, 1] [1, 1, 1, 1, 1, 1, 1] : + %0 = tensor.subtensor %arg0[0, %arg2, %arg3, 0, %arg4, 0, 0] + [1, %arg5, %arg6, 1, %arg7, 1, 1] [1, 1, 1, 1, 1, 1, 1] : tensor<1x?x?x1x?x1x1xf32> to tensor<1x?x?x1x?x1x1xf32> - %1 = subtensor %arg1[%arg2, 0, %arg3, 0, 0, %arg4, 0] - [1, %arg5, %arg6, 1, %arg7, 1, 1] [1, 1, 1, 1, 1, 1, 1] : + %1 = tensor.subtensor %arg1[%arg2, 0, %arg3, 0, 0, %arg4, 0] + [1, %arg5, %arg6, 1, %arg7, 1, 1] [1, 1, 1, 1, 1, 1, 1] : tensor<1x?x?x?x?x1x1xf32> to tensor<1x?x?x1x?x1x1xf32> return %0, %1 : tensor<1x?x?x1x?x1x1xf32>, tensor<1x?x?x1x?x1x1xf32> } // CHECK: func @fold_subtensor // CHECK-SAME: %[[ARG0:.+]]: tensor<1x?x?x1x?x1x1xf32> // CHECK-SAME: %[[ARG1:.+]]: tensor<1x?x?x?x?x1x1xf32> -// CHECK: %[[SUBTENSOR1:.+]] = subtensor %[[ARG0]] +// CHECK: %[[SUBTENSOR1:.+]] = tensor.subtensor %[[ARG0]] // CHECK-SAME: to tensor // CHECK: %[[RESULT1:.+]] = linalg.tensor_expand_shape %[[SUBTENSOR1]] // CHECK-SAME: [0, 1], [2], [3, 4, 5, 6] -// CHECK: %[[SUBTENSOR2:.+]] = subtensor %[[ARG1]] +// CHECK: %[[SUBTENSOR2:.+]] = tensor.subtensor %[[ARG1]] // CHECK-SAME: to tensor // CHECK: %[[RESULT2:.+]] = linalg.tensor_expand_shape %[[SUBTENSOR2]] // CHECK-SAME: [0, 1], [2], [3, 4, 5, 6] @@ -431,11 +431,11 @@ // ----- func @subtensor_unit_dims(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> { - %0 = subtensor %arg0[0, 2] [1, 1] [1, 1] : tensor<1x3xf32> to tensor<1x1xf32> + %0 = tensor.subtensor %arg0[0, 2] [1, 1] [1, 1] : tensor<1x3xf32> to tensor<1x1xf32> return %0 : tensor<1x1xf32> } // CHECK-LABEL: func @subtensor_unit_dims -// CHECK: %[[SUBTENSOR:.+]] = subtensor +// CHECK: %[[SUBTENSOR:.+]] = tensor.subtensor // CHECK-SAME: tensor<1x3xf32> to tensor // CHECK: %[[RESULT:.+]] = linalg.tensor_expand_shape %[[SUBTENSOR]] [] // CHECK: return %[[RESULT]] @@ -443,12 +443,12 @@ // ----- func @subtensor_insert_unit_dims(%arg0: tensor<1x3xf32>, %arg1: tensor<1x1xf32>) -> tensor<1x3xf32> { - %0 = subtensor_insert %arg1 into %arg0[0, 2] [1, 1] [1, 1] : tensor<1x1xf32> into tensor<1x3xf32> + %0 = tensor.subtensor_insert %arg1 into %arg0[0, 2] [1, 1] [1, 1] : tensor<1x1xf32> into tensor<1x3xf32> return %0 : tensor<1x3xf32> } // CHECK-LABEL: func @subtensor_insert_unit_dims // CHECK: %[[RESHAPE:.+]] = linalg.tensor_collapse_shape %{{.+}} [] -// CHECK: %[[RESULT:.+]] = subtensor_insert %[[RESHAPE]] +// CHECK: %[[RESULT:.+]] = tensor.subtensor_insert %[[RESHAPE]] // CHECK-SAME: tensor into tensor<1x3xf32> // CHECK: return %[[RESULT]] diff --git a/mlir/test/Dialect/Linalg/fusion-sequence.mlir b/mlir/test/Dialect/Linalg/fusion-sequence.mlir --- a/mlir/test/Dialect/Linalg/fusion-sequence.mlir +++ b/mlir/test/Dialect/Linalg/fusion-sequence.mlir @@ -175,18 +175,18 @@ // CHECK: %[[INIT:.+]] = linalg.init_tensor // CHECK: %[[R0:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[ARG5:.+]] = %[[INIT]]) -> (tensor) { // CHECK: %[[R1:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[ARG7:.+]] = %[[ARG5]]) -> (tensor) { -// CHECK-DAG: %[[STARG3:.+]] = subtensor %[[ARG3]] -// CHECK-DAG: %[[STARG7:.+]] = subtensor %[[ARG7]] -// CHECK-DAG: %[[STARG0:.+]] = subtensor %[[ARG0]] -// CHECK-DAG: %[[STARG1:.+]] = subtensor %[[ARG1]] -// CHECK-DAG: %[[STARG2:.+]] = subtensor %[[ARG2]] +// CHECK-DAG: %[[STARG3:.+]] = tensor.subtensor %[[ARG3]] +// CHECK-DAG: %[[STARG7:.+]] = tensor.subtensor %[[ARG7]] +// CHECK-DAG: %[[STARG0:.+]] = tensor.subtensor %[[ARG0]] +// CHECK-DAG: %[[STARG1:.+]] = tensor.subtensor %[[ARG1]] +// CHECK-DAG: %[[STARG2:.+]] = tensor.subtensor %[[ARG2]] // CHECK: %[[T0:.+]] = linalg.matmul // CHECK-SAME: ins(%[[STARG0]], %[[STARG1]] : tensor, tensor) // CHECK-SAME: outs(%[[STARG2]] : tensor) -> tensor // CHECK: %[[T1:.+]] = linalg.generic // CHECK-SAME: ins(%[[T0:.+]], %[[STARG3]] : tensor, tensor) // CHECK-SAME: outs(%[[STARG7]] : tensor) -// CHECK: %[[RESULT:.+]] = subtensor_insert %[[T1]] into %[[ARG7]] +// CHECK: %[[RESULT:.+]] = tensor.subtensor_insert %[[T1]] into %[[ARG7]] // CHECK: scf.yield %[[RESULT]] // CHECK: } // CHECK: scf.yield %[[R1]] @@ -229,21 +229,21 @@ // CHECK: %[[M_1:.+]] = memref.dim %[[ARG8]], %[[C0]] // CHECK: %[[TILE_M_1:.+]] = affine.min #[[MAP0]](%[[M_1]], %[[IV0]]) // CHECK: %[[N3:.+]] = memref.dim %[[ARG8]], %[[C1]] -// CHECK: %[[STARG6:.+]] = subtensor %[[ARG8]][%[[IV0]], 0] +// CHECK: %[[STARG6:.+]] = tensor.subtensor %[[ARG8]][%[[IV0]], 0] // CHECK-SAME: [%[[TILE_M_1]], %[[N3]]] // CHECK: %[[M_2:.+]] = memref.dim %[[ARG4]], %[[C0]] // CHECK: %[[TILE_M_2:.+]] = affine.min #[[MAP1]](%[[IV0]])[%[[M_2]], %[[M]]] // CHECK: %[[N2:.+]] = memref.dim %[[ARG4]], %[[C1]] -// CHECK: %[[STARG4:.+]] = subtensor %[[ARG4]][%[[IV0]], 0] +// CHECK: %[[STARG4:.+]] = tensor.subtensor %[[ARG4]][%[[IV0]], 0] // CHECK-SAME: [%[[TILE_M_2]], %[[N2]]] // CHECK: %[[TILE_M_3:.+]] = affine.min #[[MAP1]](%[[IV0]])[%[[M]], %[[M]]] // CHECK: %[[N0:.+]] = memref.dim %[[ARG0]], %[[C1]] -// CHECK: %[[STARG0:.+]] = subtensor %[[ARG0]][%[[IV0]], 0] +// CHECK: %[[STARG0:.+]] = tensor.subtensor %[[ARG0]][%[[IV0]], 0] // CHECK-SAME: [%[[TILE_M_3]], %[[N0]]] // CHECK: %[[M_3:.+]] = memref.dim %[[ARG2]], %[[C0]] // CHECK: %[[TILE_M_4:.+]] = affine.min #[[MAP1]](%[[IV0]])[%[[M_3]], %[[M]]] // CHECK: %[[N1:.+]] = memref.dim %[[ARG2]], %[[C1]] -// CHECK: %[[STARG2:.+]] = subtensor %[[ARG2]][%[[IV0]], 0] +// CHECK: %[[STARG2:.+]] = tensor.subtensor %[[ARG2]][%[[IV0]], 0] // CHECK-SAME: [%[[TILE_M_4]], %[[N1]]] // CHECK: %[[T0:.+]] = linalg.matmul // CHECK-SAME: ins(%[[STARG0]], %[[ARG1]] : tensor, tensor @@ -254,7 +254,7 @@ // CHECK: %[[T2:.+]] = linalg.matmul // CHECK-SAME: ins(%[[T1]], %arg5 : tensor, tensor // CHECK-SAME: ) outs(%[[STARG6]] : tensor) -// CHECK: %[[R1:.+]] = subtensor_insert %[[T2]] +// CHECK: %[[R1:.+]] = tensor.subtensor_insert %[[T2]] // CHECK-SAME: into %[[ARG8]][%[[IV0]], 0] [%[[TILE_M_1]], %[[N3]]] // CHECK: scf.yield %[[R1]] : tensor // CHECK: } diff --git a/mlir/test/Dialect/Linalg/fusion-tensor-pattern.mlir b/mlir/test/Dialect/Linalg/fusion-tensor-pattern.mlir --- a/mlir/test/Dialect/Linalg/fusion-tensor-pattern.mlir +++ b/mlir/test/Dialect/Linalg/fusion-tensor-pattern.mlir @@ -38,16 +38,16 @@ // CHECK: %[[M_2:.+]] = memref.dim %[[ARG6]], %[[C0]] // CHECK: %[[TILE_M_2:.+]] = affine.min #[[MAP1]](%[[M_2]], %[[IV0]]) // CHECK: %[[N3:.+]] = memref.dim %[[ARG6]], %[[C1]] -// CHECK: %[[ST_ARG6:.+]] = subtensor %[[ARG6]][%[[IV0]], 0] +// CHECK: %[[ST_ARG6:.+]] = tensor.subtensor %[[ARG6]][%[[IV0]], 0] // CHECK-SAME: [%[[TILE_M_2]], %[[N3]]] // CHECK: %[[TILE_M_3:.+]] = affine.min #[[MAP5]](%[[IV0]])[%[[M]], %[[M]]] // CHECK: %[[N1:.+]] = memref.dim %[[ARG0]], %[[C1]] -// CHECK: %[[ST_ARG0:.+]] = subtensor %[[ARG0]][%[[IV0]], 0] +// CHECK: %[[ST_ARG0:.+]] = tensor.subtensor %[[ARG0]][%[[IV0]], 0] // CHECK-SAME: [%[[TILE_M_3]], %[[N1]]] // CHECK: %[[M_3:.+]] = memref.dim %[[ARG2]], %[[C0]] // CHECK: %[[TILE_M_4:.+]] = affine.min #[[MAP5]](%[[IV0]])[%[[M_3]], %[[M]]] // CHECK: %[[N2_2:.+]] = memref.dim %[[ARG2]], %[[C1]] -// CHECK: %[[ST_ARG2:.+]] = subtensor %[[ARG2]][%[[IV0]], 0] +// CHECK: %[[ST_ARG2:.+]] = tensor.subtensor %[[ARG2]][%[[IV0]], 0] // CHECK-SAME: [%[[TILE_M_4]], %[[N2_2]]] // CHECK: %[[LHS:.+]] = linalg.matmul // CHECK-SAME: __internal_linalg_transform__ = "after_lhs_fusion_producer" @@ -62,30 +62,30 @@ // CHECK-SAME: %[[C0]] to %[[N2]] step %[[C16]] // CHECK-SAME: iter_args(%[[ARG10:.+]] = %[[ARG8]]) -> (tensor) { // CHECK: %[[TILE_N2:.+]] = affine.min #[[MAP2]](%[[IV2]])[%[[N2]]] -// CHECK: %[[ST_LHS:.+]] = subtensor %[[LHS]][0, %[[IV2]]] +// CHECK: %[[ST_LHS:.+]] = tensor.subtensor %[[LHS]][0, %[[IV2]]] // CHECK-SAME: [%[[TILE_M_3]], %[[TILE_N2]]] // CHECK: %[[N2_3:.+]] = memref.dim %[[ARG3]], %[[C0]] // CHECK: %[[TILE_N2_2:.+]] = affine.min #[[MAP2]](%[[IV2]])[%[[N2_3]]] // CHECK: %[[TILE_N3:.+]] = affine.min #[[MAP3]](%[[IV1]])[%[[N3_2]]] -// CHECK: %[[ST_ARG3:.+]] = subtensor %[[ARG3]][%[[IV2]], %[[IV1]]] +// CHECK: %[[ST_ARG3:.+]] = tensor.subtensor %[[ARG3]][%[[IV2]], %[[IV1]]] // CHECK-SAME: [%[[TILE_N2_2]], %[[TILE_N3]]] // CHECK: %[[M_4:.+]] = memref.dim %[[ARG10]], %[[C0]] // CHECK: %[[N3_3:.+]] = memref.dim %[[ARG10]], %[[C1]] // CHECK: %[[TILE_N3_2:.+]] = affine.min #[[MAP4]](%[[N3_3]], %[[IV1]]) -// CHECK: %[[ST_ARG4:.+]] = subtensor %[[ARG10]][0, %[[IV1]]] +// CHECK: %[[ST_ARG4:.+]] = tensor.subtensor %[[ARG10]][0, %[[IV1]]] // CHECK-SAME: [%[[M_4]], %[[TILE_N3_2]]] // CHECK: %[[ST_RESULT:.+]] = linalg.matmul // CHECK-SAME: __internal_linalg_transform__ = "after_lhs_fusion" // CHECK-SAME: ins(%[[ST_LHS]], %[[ST_ARG3]] // CHECK-SAME: : tensor, tensor) // CHECK-SAME: outs(%[[ST_ARG4]] : tensor) -// CHECK: %[[UPDATE1:.+]] = subtensor_insert %[[ST_RESULT]] +// CHECK: %[[UPDATE1:.+]] = tensor.subtensor_insert %[[ST_RESULT]] // CHECK-SAME: into %[[ARG10]][0, %[[IV1]]] [%[[M_4]], %[[TILE_N3_2]]] // CHECK: scf.yield %[[UPDATE1]] // CHECK: } // CHECK: scf.yield %[[YIELD1]] // CHECK: } -// CHECK: %[[UPDATE0:.+]] = subtensor_insert %[[YIELD0]] into +// CHECK: %[[UPDATE0:.+]] = tensor.subtensor_insert %[[YIELD0]] into // CHECK-SAME: %[[ARG6]][%[[IV0]], 0] [%[[TILE_M_2]], %[[N3]]] // CHECK: scf.yield %[[UPDATE0]] // CHECK: } @@ -114,9 +114,9 @@ // TLOOP-SAME: %[[AB_INIT_:.*]] = %[[AB_INIT]]: tensor) // TLOOP-SAME: outs (%[[ABC_INIT_:.*]] = %[[ABC_INIT]]: tensor) { -// TLOOP: %[[ABC_INIT_SUB:.*]] = subtensor %[[ABC_INIT_]][%[[IV0]], 0] -// TLOOP: %[[A_SUB:.*]] = subtensor %[[A_]][%[[IV0]], 0] -// TLOOP: %[[AB_INIT_SUB:.*]] = subtensor %[[AB_INIT_]][%[[IV0]], 0] +// TLOOP: %[[ABC_INIT_SUB:.*]] = tensor.subtensor %[[ABC_INIT_]][%[[IV0]], 0] +// TLOOP: %[[A_SUB:.*]] = tensor.subtensor %[[A_]][%[[IV0]], 0] +// TLOOP: %[[AB_INIT_SUB:.*]] = tensor.subtensor %[[AB_INIT_]][%[[IV0]], 0] // TLOOP: %[[AB_SUB:.*]] = linalg.matmul // TLOOP-SAME: ins(%[[A_SUB]], %[[B_]] : {{.*}}) outs(%[[AB_INIT_SUB]] @@ -132,19 +132,19 @@ // TLOOP-SAME: outs (%[[ABC_INIT_SUB_:.*]] = %[[ABC_INIT_SUB]]: [[TY]]) // TLOOP-SAME: iterators["parallel", "reduction"] { -// TLOOP: %[[AB_SUB_SUB:.*]] = subtensor %[[AB_SUB_]][0, %[[IV2]]] -// TLOOP: %[[C__SUB:.*]] = subtensor %[[C__]][%[[IV2]], %[[IV1]]] -// TLOOP: %[[ABS_INIT_SUB_SUB:.*]] = subtensor %[[ABC_INIT_SUB_]][0, %[[IV1]]] +// TLOOP: %[[AB_SUB_SUB:.*]] = tensor.subtensor %[[AB_SUB_]][0, %[[IV2]]] +// TLOOP: %[[C__SUB:.*]] = tensor.subtensor %[[C__]][%[[IV2]], %[[IV1]]] +// TLOOP: %[[ABS_INIT_SUB_SUB:.*]] = tensor.subtensor %[[ABC_INIT_SUB_]][0, %[[IV1]]] // TLOOP: %[[ABC_SUB_SUB:.*]] = linalg.matmul // TLOOP-SAME: ins(%[[AB_SUB_SUB]], %[[C__SUB]] : [[TY]], [[TY]]) // TLOOP-SAME: outs(%[[ABS_INIT_SUB_SUB]] : [[TY]]) -> [[TY]] -// TLOOP: %[[RES0:.*]] = subtensor_insert %[[ABC_SUB_SUB]] +// TLOOP: %[[RES0:.*]] = tensor.subtensor_insert %[[ABC_SUB_SUB]] // TLOOP-SAME: into %[[ABC_INIT_SUB_]][0, %[[IV1]]] // TLOOP: linalg.yield %[[RES0]] : [[TY]] // TLOOP: } -// TLOOP: %[[RES1:.*]] = subtensor_insert %[[ABC_SUB_]] into %[[ABC_INIT_]][%[[IV0]], 0] +// TLOOP: %[[RES1:.*]] = tensor.subtensor_insert %[[ABC_SUB_]] into %[[ABC_INIT_]][%[[IV0]], 0] // TLOOP: linalg.yield %[[RES1]] : [[TY]] // TLOOP: } // TLOOP: return %[[ABC]] : [[TY]] @@ -186,10 +186,10 @@ // CHECK-SAME: iter_args(%[[ARG4:.+]] = %{{[a-zA-Z0-9_]+}}) // CHECK: %[[YIELD:.+]] = scf.for %[[IV1:[a-zA-Z0-9_]+]] // CHECK-SAME: iter_args(%[[ARG6:.+]] = %[[ARG4]]) -// CHECK: %[[ST_ARG6:.+]] = subtensor %[[ARG6]][%[[IV0]], %[[IV1]]] -// CHECK: %[[ST_ARG0:.+]] = subtensor %[[ARG0]][%[[IV0]], 0] -// CHECK: %[[ST_ARG1:.+]] = subtensor %[[ARG1]][0, %[[IV1]]] -// CHECK: %[[ST_ARG2:.+]] = subtensor %[[ARG2]][%[[IV0]], %[[IV1]]] +// CHECK: %[[ST_ARG6:.+]] = tensor.subtensor %[[ARG6]][%[[IV0]], %[[IV1]]] +// CHECK: %[[ST_ARG0:.+]] = tensor.subtensor %[[ARG0]][%[[IV0]], 0] +// CHECK: %[[ST_ARG1:.+]] = tensor.subtensor %[[ARG1]][0, %[[IV1]]] +// CHECK: %[[ST_ARG2:.+]] = tensor.subtensor %[[ARG2]][%[[IV0]], %[[IV1]]] // CHECK: %[[LHS:.+]] = linalg.matmul // CHECK-SAME: ins(%[[ST_ARG0]], %[[ST_ARG1]] // CHECK-SAME: : tensor, tensor) @@ -197,7 +197,7 @@ // CHECK: %[[ST_RESULT:.+]] = linalg.generic // CHECK-SAME: ins(%[[LHS]] : tensor) // CHECK-SAME: outs(%[[ST_ARG6]] : tensor) -// CHECK: %[[UPDATE:.+]] = subtensor_insert %[[ST_RESULT]] +// CHECK: %[[UPDATE:.+]] = tensor.subtensor_insert %[[ST_RESULT]] // CHECK-SAME: into %[[ARG6]][%[[IV0]], %[[IV1]]] // CHECK: scf.yield %[[UPDATE]] // CHECK: scf.yield %[[YIELD]] @@ -226,10 +226,10 @@ // TLOOP-SAME: %[[AB_:.*]] = %[[AB]]: [[TY]]) // TLOOP-SAME: outs (%[[INIT_:.*]] = %[[INIT]]: [[TY]]) { -// TLOOP: %[[INIT_SUB:.*]] = subtensor %[[INIT_]][%[[IV0]], %[[IV1]]] -// TLOOP: %[[A_SUB:.*]] = subtensor %[[A_]][%[[IV0]], 0] -// TLOOP: %[[B_SUB:.*]] = subtensor %[[B_]][0, %[[IV1]]] -// TLOOP: %[[AB_SUB_INIT:.*]] = subtensor %[[AB_]][%[[IV0]], %[[IV1]]] +// TLOOP: %[[INIT_SUB:.*]] = tensor.subtensor %[[INIT_]][%[[IV0]], %[[IV1]]] +// TLOOP: %[[A_SUB:.*]] = tensor.subtensor %[[A_]][%[[IV0]], 0] +// TLOOP: %[[B_SUB:.*]] = tensor.subtensor %[[B_]][0, %[[IV1]]] +// TLOOP: %[[AB_SUB_INIT:.*]] = tensor.subtensor %[[AB_]][%[[IV0]], %[[IV1]]] // TLOOP: %[[AB_SUB:.*]] = linalg.matmul // TLOOP-SAME: ins(%[[A_SUB]], %[[B_SUB]] : [[TY]], [[TY]]) @@ -238,7 +238,7 @@ // TLOOP: %[[DOUBLE_AB:.*]] = linalg.generic // TLOOP-SAME: ins(%[[AB_SUB]] : [[TY]]) outs(%[[INIT_SUB]] : [[TY]]) -// TLOOP: %[[RESULT_SUB:.*]] = subtensor_insert +// TLOOP: %[[RESULT_SUB:.*]] = tensor.subtensor_insert // TLOOP-SAME: %[[DOUBLE_AB:.*]] into %[[INIT_]][%[[IV0]], %[[IV1]]] // TLOOP: linalg.yield %[[RESULT_SUB]] : [[TY]] @@ -267,13 +267,13 @@ // CHECK-NOT: fill // CHECK: scf.for %[[I:.*]]{{.*}}iter_args(%{{.*}} = %[[ARG0]]) -> (tensor) { // CHECK: scf.for %[[J:.*]] -// CHECK: %[[ST:.*]] = subtensor %[[ARG0]] +// CHECK: %[[ST:.*]] = tensor.subtensor %[[ARG0]] // CHECK: %[[ST_FILL:.*]] = linalg.fill(%[[ST]], %[[C0]]) {__internal_linalg_transform__ = "after_out_fusion_producer"} : tensor, f32 -> tensor // CHECK: %[[ST_MM_RES:.*]] = scf.for %[[K:.*]]{{.*}}iter_args(%[[BB:.*]] = %[[ST_FILL]]) -> (tensor) { // CHECK-NOT: fill // CHECK: %[[ST_MM:.*]] = linalg.matmul {__internal_linalg_transform__ = "after_out_fusion"} ins(%{{.*}}, %{{.*}} : tensor, tensor) outs(%[[BB]] : tensor) -> tensor // CHECK: scf.yield %[[ST_MM]] : tensor -// CHECK: %[[MM:.*]] = subtensor_insert %[[ST_MM_RES]] into {{.*}} +// CHECK: %[[MM:.*]] = tensor.subtensor_insert %[[ST_MM_RES]] into {{.*}} // CHECK: scf.yield %[[MM]] : tensor @@ -300,9 +300,9 @@ // TLOOP-SAME: outs (%[[OUT_:.*]] = %[[OUT]]: [[TY]]) { // TLOOP: %[[DIM_A__1:.*]] = memref.dim %[[A_]], %[[C1]] : [[TY]] -// TLOOP: %[[A_SUB:.*]] = subtensor %[[A_]][%[[I]], 0] -// TLOOP: %[[B_SUB:.*]] = subtensor %[[B_]][0, %[[J]]] -// TLOOP: %[[OUT_SUB:.*]] = subtensor %[[OUT_]][%[[I]], %[[J]]] +// TLOOP: %[[A_SUB:.*]] = tensor.subtensor %[[A_]][%[[I]], 0] +// TLOOP: %[[B_SUB:.*]] = tensor.subtensor %[[B_]][0, %[[J]]] +// TLOOP: %[[OUT_SUB:.*]] = tensor.subtensor %[[OUT_]][%[[I]], %[[J]]] // TLOOP: %[[INIT_SUB:.*]] = linalg.fill(%[[OUT_SUB]], %[[C0_F32]]) // TLOOP: %[[AB_SUB:.*]] = linalg.tiled_loop (%[[K:.*]]) = (%[[C0]]) @@ -312,15 +312,15 @@ // TLOOP-SAME: outs (%[[INIT_SUB_:.*]] = %[[INIT_SUB]]: [[TY]]) // TLOOP-SAME: iterators["reduction"] { -// TLOOP: %[[A_SUB_SUB:.*]] = subtensor %[[A_SUB_]][0, %[[K]]] -// TLOOP: %[[B_SUB_SUB:.*]] = subtensor %[[B_SUB_]][%[[K]], 0] +// TLOOP: %[[A_SUB_SUB:.*]] = tensor.subtensor %[[A_SUB_]][0, %[[K]]] +// TLOOP: %[[B_SUB_SUB:.*]] = tensor.subtensor %[[B_SUB_]][%[[K]], 0] // TLOOP: %[[AB_SUB_SUB:.*]] = linalg.matmul // TLOOP-SAME: ins(%[[A_SUB_SUB]], %[[B_SUB_SUB]] : [[TY]], [[TY]]) // TLOOP-SAME: outs(%[[INIT_SUB_]] : [[TY]]) -> [[TY]] // TLOOP: linalg.yield %[[AB_SUB_SUB]] : [[TY]] // TLOOP: } -// TLOOP: %[[SUB_RESULT:.*]] = subtensor_insert %[[AB_SUB]] +// TLOOP: %[[SUB_RESULT:.*]] = tensor.subtensor_insert %[[AB_SUB]] // TLOOP-SAME: into %[[OUT_]][%[[I]], %[[J]]] // TLOOP: linalg.yield %[[SUB_RESULT]] : [[TY]] // TLOOP: } @@ -371,9 +371,9 @@ // TLOOP-SAME: outs (%[[OUT_:.*]] = %[[OUT]]: [[TY]]) { // TLOOP: %[[DIM_A__1:.*]] = memref.dim %[[A_]], %[[C1]] : [[TY]] -// TLOOP: %[[A_SUB:.*]] = subtensor %[[A_]][%[[I]], 0] -// TLOOP: %[[B_SUB:.*]] = subtensor %[[B_]][0, %[[J]]] -// TLOOP: %[[OUT_SUB:.*]] = subtensor %[[OUT_]][%[[I]], %[[J]]] +// TLOOP: %[[A_SUB:.*]] = tensor.subtensor %[[A_]][%[[I]], 0] +// TLOOP: %[[B_SUB:.*]] = tensor.subtensor %[[B_]][0, %[[J]]] +// TLOOP: %[[OUT_SUB:.*]] = tensor.subtensor %[[OUT_]][%[[I]], %[[J]]] // TLOOP: %[[INIT_SUB:.*]] = linalg.generic // TLOOP-SAME: ins(%[[C0_F32_]] // TLOOP-SAME: outs(%[[OUT_SUB]] @@ -385,15 +385,15 @@ // TLOOP-SAME: outs (%[[INIT_SUB_:.*]] = %[[INIT_SUB]]: [[TY]]) // TLOOP-SAME: iterators["reduction"] { -// TLOOP: %[[A_SUB_SUB:.*]] = subtensor %[[A_SUB_]][0, %[[K]]] -// TLOOP: %[[B_SUB_SUB:.*]] = subtensor %[[B_SUB_]][%[[K]], 0] +// TLOOP: %[[A_SUB_SUB:.*]] = tensor.subtensor %[[A_SUB_]][0, %[[K]]] +// TLOOP: %[[B_SUB_SUB:.*]] = tensor.subtensor %[[B_SUB_]][%[[K]], 0] // TLOOP: %[[AB_SUB_SUB:.*]] = linalg.matmul // TLOOP-SAME: ins(%[[A_SUB_SUB]], %[[B_SUB_SUB]] : [[TY]], [[TY]]) // TLOOP-SAME: outs(%[[INIT_SUB_]] : [[TY]]) -> [[TY]] // TLOOP: linalg.yield %[[AB_SUB_SUB]] : [[TY]] // TLOOP: } -// TLOOP: %[[SUB_RESULT:.*]] = subtensor_insert %[[AB_SUB]] +// TLOOP: %[[SUB_RESULT:.*]] = tensor.subtensor_insert %[[AB_SUB]] // TLOOP-SAME: into %[[OUT_]][%[[I]], %[[J]]] // TLOOP: linalg.yield %[[SUB_RESULT]] : [[TY]] // TLOOP: } diff --git a/mlir/test/Dialect/Linalg/hoist-padding.mlir b/mlir/test/Dialect/Linalg/hoist-padding.mlir --- a/mlir/test/Dialect/Linalg/hoist-padding.mlir +++ b/mlir/test/Dialect/Linalg/hoist-padding.mlir @@ -53,10 +53,10 @@ // CHECK: %[[A:.*]] = scf.for %[[J1:[0-9a-z]+]] = // Iteration count along J1 // CHECK: %[[IDXpad0_K:[0-9]+]] = affine.apply #[[$DIV4]](%[[J1]]) - // CHECK: subtensor %{{.*}} [1, 1] : tensor to tensor + // CHECK: tensor.subtensor %{{.*}} [1, 1] : tensor to tensor // CHECK: linalg.pad_tensor %{{.*}} // CHECK: : tensor to tensor<2x4xf32> - // CHECK: subtensor_insert %{{.*}} into %{{.*}}[%[[IDXpad0_K]], 0, 0] + // CHECK: tensor.subtensor_insert %{{.*}} into %{{.*}}[%[[IDXpad0_K]], 0, 0] // CHECK-SAME: [1, 2, 4] [1, 1, 1] : tensor<2x4xf32> into tensor // Second tensor is KxN but loop order is (M, N, K) so padded tensor is NxKx4x3 // CHECK: %[[SZpad1_N:[0-9]+]] = affine.apply #[[$DIVS3]]()[%[[dN]]] @@ -69,23 +69,23 @@ // CHECK: scf.for %[[J2:[0-9a-z]+]] = // Iteration count along J2 // CHECK: %[[IDXpad1_N:[0-9]+]] = affine.apply #[[$DIV4]](%[[J2]]) - // CHECK: subtensor %{{.*}} [1, 1] : tensor to tensor + // CHECK: tensor.subtensor %{{.*}} [1, 1] : tensor to tensor // CHECK: linalg.pad_tensor %{{.*}} // CHECK: : tensor to tensor<4x3xf32> - // CHECK: subtensor_insert %{{.*}} into %{{.*}}[%[[IDXpad1_K]], %[[IDXpad1_N]], 0, 0] + // CHECK: tensor.subtensor_insert %{{.*}} into %{{.*}}[%[[IDXpad1_K]], %[[IDXpad1_N]], 0, 0] // CHECK-SAME: [1, 1, 4, 3] [1, 1, 1, 1] : tensor<4x3xf32> into tensor // 2-D loop // CHECK: scf.for %[[J:[0-9a-zA-Z]+]] // CHECK: scf.for %[[K:[0-9a-zA-Z]+]] // Iteration count along K // CHECK: %[[IDXpad0_K:[0-9]+]] = affine.apply #[[$DIV4]](%[[K]]) - // CHECK: %[[stA:.*]] = subtensor %[[A]][%[[IDXpad0_K]], 0, 0] [1, 2, 4] [1, 1, 1] : + // CHECK: %[[stA:.*]] = tensor.subtensor %[[A]][%[[IDXpad0_K]], 0, 0] [1, 2, 4] [1, 1, 1] : // CHECK-SAME: tensor to tensor<2x4xf32> // Iteration count along J // CHECK: %[[IDXpad1_N:[0-9]+]] = affine.apply #[[$DIV3]](%[[J]]) // Iteration count along K // CHECK: %[[IDXpad1_K:[0-9]+]] = affine.apply #[[$DIV4]](%[[K]]) - // CHECK: %[[stB:.*]] = subtensor %[[B]][%[[IDXpad1_N]], %[[IDXpad1_K]], 0, 0] [1, 1, 4, 3] [1, 1, 1, 1] : + // CHECK: %[[stB:.*]] = tensor.subtensor %[[B]][%[[IDXpad1_N]], %[[IDXpad1_K]], 0, 0] [1, 1, 4, 3] [1, 1, 1, 1] : // CHECK-SAME: tensor to tensor<4x3xf32> // CHECK: %[[stC:.*]] = linalg.pad_tensor %{{.*}} // CHECK: : tensor to tensor<2x3xf32> @@ -98,17 +98,17 @@ %7 = affine.min #map0(%arg3)[%6] %8 = memref.dim %arg0, %c1 : tensor %9 = affine.min #map1(%arg7)[%8] - %10 = subtensor %arg0[%arg3, %arg7] [%7, %9] [1, 1] : tensor to tensor + %10 = tensor.subtensor %arg0[%arg3, %arg7] [%7, %9] [1, 1] : tensor to tensor %11 = memref.dim %arg1, %c0 : tensor %12 = affine.min #map1(%arg7)[%11] %13 = memref.dim %arg1, %c1 : tensor %14 = affine.min #map2(%arg5)[%13] - %15 = subtensor %arg1[%arg7, %arg5] [%12, %14] [1, 1] : tensor to tensor + %15 = tensor.subtensor %arg1[%arg7, %arg5] [%12, %14] [1, 1] : tensor to tensor %16 = memref.dim %arg8, %c0 : tensor %17 = affine.min #map3(%16, %arg3) %18 = memref.dim %arg8, %c1 : tensor %19 = affine.min #map4(%18, %arg5) - %20 = subtensor %arg8[%arg3, %arg5] [%17, %19] [1, 1] : tensor to tensor + %20 = tensor.subtensor %arg8[%arg3, %arg5] [%17, %19] [1, 1] : tensor to tensor %21 = subi %c2, %7 : index %22 = subi %c4, %9 : index %23 = linalg.pad_tensor %10 low[%c0, %c0] high[%21, %22] { @@ -128,8 +128,8 @@ linalg.yield %cst : f32 } : tensor to tensor<2x3xf32> %30 = linalg.matmul ins(%23, %26 : tensor<2x4xf32>, tensor<4x3xf32>) outs(%29 : tensor<2x3xf32>) -> tensor<2x3xf32> - %31 = subtensor %30[0, 0] [%7, %14] [1, 1] : tensor<2x3xf32> to tensor - %32 = subtensor_insert %31 into %arg8[%arg3, %arg5] [%17, %19] [%c1, %c1] : tensor into tensor + %31 = tensor.subtensor %30[0, 0] [%7, %14] [1, 1] : tensor<2x3xf32> to tensor + %32 = tensor.subtensor_insert %31 into %arg8[%arg3, %arg5] [%17, %19] [%c1, %c1] : tensor into tensor scf.yield %32 : tensor } scf.yield %5 : tensor @@ -173,7 +173,7 @@ // CHECK: %[[INIT_PACKED_A:.*]] = linalg.init_tensor [%[[D0]], %[[D1]], 2] : tensor // CHECK: %[[PACKED_A:.*]] = scf.for %[[II:[0-9a-z]+]] = {{.*}} iter_args(%{{.*}} = %[[INIT_PACKED_A]]) -> (tensor) { // CHECK: scf.for %[[III:[0-9a-z]+]] = - // CHECK: subtensor_insert %{{.*}} into %{{.*}}[%{{.*}}, %{{.*}}, 0] [1, 1, 2] [1, 1, 1] : tensor<2xf32> into tensor + // CHECK: tensor.subtensor_insert %{{.*}} into %{{.*}}[%{{.*}}, %{{.*}}, 0] [1, 1, 2] [1, 1, 1] : tensor<2xf32> into tensor // // CHECK: %[[D0_2:.*]] = affine.apply #[[$DIV4]](%[[MR8]]) // CHECK: %[[MM4_2:.*]] = affine.min #[[$MIN_MOD4]](%[[MR8]]) @@ -182,33 +182,33 @@ // CHECK: %[[INIT_PACKED_B:.*]] = linalg.init_tensor [%[[D0_2]], %[[D1_2]], 2] : tensor // CHECK: %[[PACKED_B:.*]] = scf.for %[[II_2:[0-9a-z]+]] = {{.*}} iter_args(%{{.*}} = %[[INIT_PACKED_B]]) -> (tensor) { // CHECK: scf.for %[[III_2:[0-9a-z]+]] = - // CHECK: subtensor_insert %{{.*}} into %{{.*}}[%{{.*}}, %{{.*}}, 0] [1, 1, 2] [1, 1, 1] : tensor<2xf32> into tensor + // CHECK: tensor.subtensor_insert %{{.*}} into %{{.*}}[%{{.*}}, %{{.*}}, 0] [1, 1, 2] [1, 1, 1] : tensor<2xf32> into tensor // Compute. // CHECK: scf.for %[[II_3:[0-9a-z]+]] = // CHECK: scf.for %[[III_3:[0-9a-z]+]] = {{.*}} iter_args(%[[C:.*]] = %{{.*}}) -> (tensor) { // CHECK: %[[IDX0:.*]] = affine.apply #[[$DIV4]](%[[II_3]]) // CHECK: %[[IDX1:.*]] = affine.apply #[[$DIV2]](%[[III_3]]) - // CHECK: %[[A:.*]] = subtensor %[[PACKED_A]][%[[IDX0]], %[[IDX1]], 0] [1, 1, 2] [1, 1, 1] : tensor to tensor<2xf32> + // CHECK: %[[A:.*]] = tensor.subtensor %[[PACKED_A]][%[[IDX0]], %[[IDX1]], 0] [1, 1, 2] [1, 1, 1] : tensor to tensor<2xf32> // CHECK: %[[IDX0_2:.*]] = affine.apply #[[$DIV4]](%[[II_3]]) // CHECK: %[[IDX1_2:.*]] = affine.apply #[[$DIV2]](%[[III_3]]) - // CHECK: %[[B:.*]] = subtensor %[[PACKED_B]][%[[IDX0_2]], %[[IDX1_2]], 0] [1, 1, 2] [1, 1, 1] : tensor to tensor<2xf32> + // CHECK: %[[B:.*]] = tensor.subtensor %[[PACKED_B]][%[[IDX0_2]], %[[IDX1_2]], 0] [1, 1, 2] [1, 1, 1] : tensor to tensor<2xf32> // CHECK: linalg.dot ins(%[[A]], %[[B]] : tensor<2xf32>, tensor<2xf32>) outs(%[[C]] : tensor) -> tensor %4 = scf.for %arg3 = %c0 to %1 step %c8 iter_args(%arg4 = %arg2) -> (tensor) { %5 = affine.min #map0(%arg3)[%2] - %6 = subtensor %arg0[%arg3] [%5] [1] : tensor to tensor + %6 = tensor.subtensor %arg0[%arg3] [%5] [1] : tensor to tensor %7 = affine.min #map0(%arg3)[%3] - %8 = subtensor %arg1[%arg3] [%7] [1] : tensor to tensor + %8 = tensor.subtensor %arg1[%arg3] [%7] [1] : tensor to tensor %9 = scf.for %arg5 = %c0 to %5 step %c4 iter_args(%arg6 = %arg4) -> (tensor) { %10 = affine.min #map1(%5, %arg5) - %11 = subtensor %6[%arg5] [%10] [1] : tensor to tensor + %11 = tensor.subtensor %6[%arg5] [%10] [1] : tensor to tensor %12 = affine.min #map1(%7, %arg5) - %13 = subtensor %8[%arg5] [%12] [1] : tensor to tensor + %13 = tensor.subtensor %8[%arg5] [%12] [1] : tensor to tensor %14 = scf.for %arg7 = %c0 to %10 step %c2 iter_args(%arg8 = %arg6) -> (tensor) { %15 = affine.min #map2(%10, %arg7) - %16 = subtensor %11[%arg7] [%15] [1] : tensor to tensor + %16 = tensor.subtensor %11[%arg7] [%15] [1] : tensor to tensor %17 = affine.min #map2(%12, %arg7) - %18 = subtensor %13[%arg7] [%17] [1] : tensor to tensor + %18 = tensor.subtensor %13[%arg7] [%17] [1] : tensor to tensor %19 = subi %c2, %15 : index %20 = linalg.pad_tensor %16 low[%c0] high[%19] { ^bb0(%arg9: index): // no predecessors @@ -245,17 +245,17 @@ %1 = scf.for %arg3 = %c0 to %c32 step %c16 iter_args(%arg4 = %arg2) -> (tensor<32x64xf32>) { %2 = scf.for %arg5 = %c0 to %c64 step %c32 iter_args(%arg6 = %arg4) -> (tensor<32x64xf32>) { %3 = scf.for %arg7 = %c0 to %c128 step %c32 iter_args(%arg8 = %arg6) -> (tensor<32x64xf32>) { - %4 = subtensor %arg0[%arg3, %arg7] [16, 32] [1, 1] : tensor<32x128xf32> to tensor<16x32xf32> - %5 = subtensor %arg1[%arg7, %arg5] [32, 32] [1, 1] : tensor<128x64xf32> to tensor<32x32xf32> - %6 = subtensor %arg8[%arg3, %arg5] [16, 32] [1, 1] : tensor<32x64xf32> to tensor<16x32xf32> + %4 = tensor.subtensor %arg0[%arg3, %arg7] [16, 32] [1, 1] : tensor<32x128xf32> to tensor<16x32xf32> + %5 = tensor.subtensor %arg1[%arg7, %arg5] [32, 32] [1, 1] : tensor<128x64xf32> to tensor<32x32xf32> + %6 = tensor.subtensor %arg8[%arg3, %arg5] [16, 32] [1, 1] : tensor<32x64xf32> to tensor<16x32xf32> %7 = scf.for %arg9 = %c0 to %c16 step %c2 iter_args(%arg10 = %6) -> (tensor<16x32xf32>) { %10 = scf.for %arg11 = %c0 to %c32 step %c4 iter_args(%arg12 = %arg10) -> (tensor<16x32xf32>) { %11 = scf.for %arg13 = %c0 to %c32 step %c16 iter_args(%arg14 = %arg12) -> (tensor<16x32xf32>) { - %12 = subtensor %4[%arg9, %arg13] [2, 16] [1, 1] : tensor<16x32xf32> to tensor<2x16xf32> + %12 = tensor.subtensor %4[%arg9, %arg13] [2, 16] [1, 1] : tensor<16x32xf32> to tensor<2x16xf32> %13 = tensor.cast %12 : tensor<2x16xf32> to tensor - %14 = subtensor %5[%arg13, %arg11] [16, 4] [1, 1] : tensor<32x32xf32> to tensor<16x4xf32> + %14 = tensor.subtensor %5[%arg13, %arg11] [16, 4] [1, 1] : tensor<32x32xf32> to tensor<16x4xf32> %15 = tensor.cast %14 : tensor<16x4xf32> to tensor - %16 = subtensor %arg14[%arg9, %arg11] [2, 4] [1, 1] : tensor<16x32xf32> to tensor<2x4xf32> + %16 = tensor.subtensor %arg14[%arg9, %arg11] [2, 4] [1, 1] : tensor<16x32xf32> to tensor<2x4xf32> %17 = tensor.cast %16 : tensor<2x4xf32> to tensor %18 = linalg.pad_tensor %13 low[%c0, %c0] high[%c0, %c0] { ^bb0(%arg15: index, %arg16: index): // no predecessors @@ -271,7 +271,7 @@ } : tensor to tensor<2x4xf32> %21 = linalg.matmul ins(%18, %19 : tensor<2x16xf32>, tensor<16x4xf32>) outs(%20 : tensor<2x4xf32>) -> tensor<2x4xf32> %22 = tensor.cast %21 : tensor<2x4xf32> to tensor - %23 = subtensor_insert %22 into %arg14[%arg9, %arg11] [%c2, %c4] [1, 1] : tensor into tensor<16x32xf32> + %23 = tensor.subtensor_insert %22 into %arg14[%arg9, %arg11] [%c2, %c4] [1, 1] : tensor into tensor<16x32xf32> scf.yield %23 : tensor<16x32xf32> } scf.yield %11 : tensor<16x32xf32> @@ -279,7 +279,7 @@ scf.yield %10 : tensor<16x32xf32> } %8 = tensor.cast %7 : tensor<16x32xf32> to tensor - %9 = subtensor_insert %8 into %arg8[%arg3, %arg5] [%c16, %c32] [1, 1] : tensor into tensor<32x64xf32> + %9 = tensor.subtensor_insert %8 into %arg8[%arg3, %arg5] [%c16, %c32] [1, 1] : tensor into tensor<32x64xf32> scf.yield %9 : tensor<32x64xf32> } scf.yield %3 : tensor<32x64xf32> diff --git a/mlir/test/Dialect/Linalg/hoisting.mlir b/mlir/test/Dialect/Linalg/hoisting.mlir --- a/mlir/test/Dialect/Linalg/hoisting.mlir +++ b/mlir/test/Dialect/Linalg/hoisting.mlir @@ -349,7 +349,7 @@ -> (tensor, tensor, tensor) { // Hoisted - // CHECK: %[[ST0:.*]] = subtensor %[[TENSOR0_ARG]][%[[I]], %[[I]]]{{.*}}: tensor to tensor + // CHECK: %[[ST0:.*]] = tensor.subtensor %[[TENSOR0_ARG]][%[[I]], %[[I]]]{{.*}}: tensor to tensor // CHECK: %[[V0:.*]] = vector.transfer_read %[[ST0]]{{.*}} : tensor, vector<1xf32> // CHECK: %[[R:.*]]:3 = scf.for %[[J:.*]] = {{.*}} iter_args( @@ -362,19 +362,19 @@ iter_args(%arg6 = %arg0, %arg7 = %arg1, %arg8 = %arg2) -> (tensor, tensor, tensor) { // Hoists. - %st0 = subtensor %arg6[%i, %i][%step, %step][1, 1] : tensor to tensor + %st0 = tensor.subtensor %arg6[%i, %i][%step, %step][1, 1] : tensor to tensor %r0 = vector.transfer_read %st0[%c0, %c0], %cst: tensor, vector<1xf32> - // CHECK: %[[ST1:.*]] = subtensor %[[TENSOR1_ARG_L2]][%[[J]],{{.*}}: tensor to tensor + // CHECK: %[[ST1:.*]] = tensor.subtensor %[[TENSOR1_ARG_L2]][%[[J]],{{.*}}: tensor to tensor // CHECK: %[[V1:.*]] = vector.transfer_read %[[ST1]]{{.*}} : tensor, vector<2xf32> // Does not hoist (subtensor depends on %j) - %st1 = subtensor %arg7[%j, %c0][%step, %step][1, 1] : tensor to tensor + %st1 = tensor.subtensor %arg7[%j, %c0][%step, %step][1, 1] : tensor to tensor %r1 = vector.transfer_read %st1[%c0, %c0], %cst: tensor, vector<2xf32> - // CHECK: %[[ST2:.*]] = subtensor %[[TENSOR2_ARG_L2]][%[[I]],{{.*}}: tensor to tensor + // CHECK: %[[ST2:.*]] = tensor.subtensor %[[TENSOR2_ARG_L2]][%[[I]],{{.*}}: tensor to tensor // CHECK: %[[V2:.*]] = vector.transfer_read %[[ST2]]{{.*}} : tensor, vector<3xf32> // Does not hoist, 2 subtensor %arg8. - %st2 = subtensor %arg8[%i, %c0][%step, %step][1, 1] : tensor to tensor + %st2 = tensor.subtensor %arg8[%i, %c0][%step, %step][1, 1] : tensor to tensor %r2 = vector.transfer_read %st2[%c0, %c0], %cst: tensor, vector<3xf32> // CHECK: %[[U0:.*]] = "some_use"(%[[V0_ARG_L2]]) : (vector<1xf32>) -> vector<1xf32> @@ -396,17 +396,17 @@ %w2 = vector.transfer_write %u2, %st2[%c0, %c0] : vector<3xf32>, tensor // Hoists. - %sti0 = subtensor_insert %w0 into %arg6[%i, %i][%step, %step][1, 1] : tensor into tensor + %sti0 = tensor.subtensor_insert %w0 into %arg6[%i, %i][%step, %step][1, 1] : tensor into tensor - // CHECK-DAG: subtensor_insert %[[STI1]] into %[[TENSOR1_ARG_L2]][%[[J]],{{.*}}: tensor into tensor + // CHECK-DAG: tensor.subtensor_insert %[[STI1]] into %[[TENSOR1_ARG_L2]][%[[J]],{{.*}}: tensor into tensor // Does not hoist (depends on %j). - %sti1 = subtensor_insert %w1 into %arg7[%j, %c0][%step, %step][1, 1] : tensor into tensor + %sti1 = tensor.subtensor_insert %w1 into %arg7[%j, %c0][%step, %step][1, 1] : tensor into tensor - // CHECK-DAG: subtensor_insert %[[STI2]] into %[[TENSOR2_ARG_L2]][%[[I]],{{.*}}: tensor into tensor + // CHECK-DAG: tensor.subtensor_insert %[[STI2]] into %[[TENSOR2_ARG_L2]][%[[I]],{{.*}}: tensor into tensor // Does not hoist, 2 subtensor / subtensor_insert for %arg8. - %sti2 = subtensor_insert %w2 into %arg8[%i, %c0][%step, %step][1, 1] : tensor into tensor - %st22 = subtensor %sti2[%i, %c0][%step, %step][1, 1] : tensor to tensor - %sti22 = subtensor_insert %st22 into %arg8[%i, %c0][%step, %step][1, 1] : tensor into tensor + %sti2 = tensor.subtensor_insert %w2 into %arg8[%i, %c0][%step, %step][1, 1] : tensor into tensor + %st22 = tensor.subtensor %sti2[%i, %c0][%step, %step][1, 1] : tensor to tensor + %sti22 = tensor.subtensor_insert %st22 into %arg8[%i, %c0][%step, %step][1, 1] : tensor into tensor // CHECK: scf.yield {{.*}} : tensor, tensor, vector<1xf32> // CHECK: } @@ -416,7 +416,7 @@ // Hoisted // CHECK: %[[STI0:.*]] = vector.transfer_write %[[R]]#2, %[[ST0]]{{.*}} : vector<1xf32>, tensor - // CHECK: subtensor_insert %[[STI0]] into %[[TENSOR0_ARG]][%[[I]], %[[I]]]{{.*}} : tensor into tensor + // CHECK: tensor.subtensor_insert %[[STI0]] into %[[TENSOR0_ARG]][%[[I]], %[[I]]]{{.*}} : tensor into tensor // CHECK: scf.yield {{.*}} : tensor, tensor, tensor scf.yield %1#0, %1#1, %1#2 : diff --git a/mlir/test/Dialect/Linalg/roundtrip.mlir b/mlir/test/Dialect/Linalg/roundtrip.mlir --- a/mlir/test/Dialect/Linalg/roundtrip.mlir +++ b/mlir/test/Dialect/Linalg/roundtrip.mlir @@ -732,11 +732,11 @@ %prod = linalg.tiled_loop (%i) = (%c0) to (%c24) step (%c4) ins(%lhs_ = %lhs: tensor<24x64xi8>, %rhs_ = %rhs: tensor<24x64xi8>) outs(%out_ = %out: tensor<24x64xi8>) { - %lhs_sub = subtensor %lhs_[%i, 0] [%c4, %c64] [1, 1] + %lhs_sub = tensor.subtensor %lhs_[%i, 0] [%c4, %c64] [1, 1] : tensor<24x64xi8> to tensor - %rhs_sub = subtensor %rhs_[%i, 0] [%c4, %c64] [1, 1] + %rhs_sub = tensor.subtensor %rhs_[%i, 0] [%c4, %c64] [1, 1] : tensor<24x64xi8> to tensor - %out_sub = subtensor %out_[%i, 0] [%c4, %c64] [1, 1] + %out_sub = tensor.subtensor %out_[%i, 0] [%c4, %c64] [1, 1] : tensor<24x64xi8> to tensor %sum = linalg.generic #trait_4 @@ -747,7 +747,7 @@ linalg.yield %s : i8 } -> tensor - %sum_sub = subtensor_insert %sum into %out_[%i, 0][%c4, %c64][1, 1] + %sum_sub = tensor.subtensor_insert %sum into %out_[%i, 0][%c4, %c64][1, 1] : tensor into tensor<24x64xi8> linalg.yield %sum_sub : tensor<24x64xi8> } @@ -792,13 +792,13 @@ outs(%o_ = %output: tensor<24xf32>) iterators["reduction", "parallel", "reduction"] distribution["block_x", "block_y", "none"] { - %sub_3d = subtensor %i3d_[%i, %j, %k][2, 4, 8][1, 1, 1] + %sub_3d = tensor.subtensor %i3d_[%i, %j, %k][2, 4, 8][1, 1, 1] : tensor<16x24x32xf32> to tensor<2x4x8xf32> - %sub_2d = subtensor %i2d_[%i, %k][2, 8][1, 1] + %sub_2d = tensor.subtensor %i2d_[%i, %k][2, 8][1, 1] : tensor<16x32xf32> to tensor<2x8xf32> - %sub_1d = subtensor %i1d_[%j] [4] [1] + %sub_1d = tensor.subtensor %i1d_[%j] [4] [1] : tensor<24xf32> to tensor<4xf32> - %sub_out = subtensor %o_[%j] [4] [1] + %sub_out = tensor.subtensor %o_[%j] [4] [1] : tensor<24xf32> to tensor<4xf32> %acc = linalg.generic #trait_5 ins(%sub_3d, %sub_2d, %sub_1d @@ -810,7 +810,7 @@ linalg.yield %1 : f32 } -> tensor<4xf32> - %sum_sub = subtensor_insert %acc into %o_[%j][%c4][1] + %sum_sub = tensor.subtensor_insert %acc into %o_[%j][%c4][1] : tensor<4xf32> into tensor<24xf32> linalg.yield %sum_sub : tensor<24xf32> } diff --git a/mlir/test/Dialect/Linalg/tile-and-distribute.mlir b/mlir/test/Dialect/Linalg/tile-and-distribute.mlir --- a/mlir/test/Dialect/Linalg/tile-and-distribute.mlir +++ b/mlir/test/Dialect/Linalg/tile-and-distribute.mlir @@ -199,12 +199,12 @@ // CHECK: %[[STEPX:.+]] = affine.apply #[[MULMAP]]()[%[[NBLOCKSX]], %[[C8]]] // CHECK: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor) { // CHECK: %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor) { -// CHECK: %[[sTA:.*]] = subtensor %[[TA]][{{.*}}] : tensor to tensor -// CHECK: %[[sTB:.*]] = subtensor %[[TB]][{{.*}}] : tensor to tensor -// CHECK: %[[sTC:.*]] = subtensor %[[TC2]][{{.*}}] : tensor to tensor +// CHECK: %[[sTA:.*]] = tensor.subtensor %[[TA]][{{.*}}] : tensor to tensor +// CHECK: %[[sTB:.*]] = tensor.subtensor %[[TB]][{{.*}}] : tensor to tensor +// CHECK: %[[sTC:.*]] = tensor.subtensor %[[TC2]][{{.*}}] : tensor to tensor // CHECK: %[[sTD:.*]] = linalg.matmul ins(%[[sTA]], %[[sTB]] : tensor, tensor) // CHECK-SAME: outs(%[[sTC]] : tensor) -> tensor -// CHECK: %[[TD:.*]] = subtensor_insert %[[sTD]] into %[[TC2]][{{.*}}] : tensor into tensor +// CHECK: %[[TD:.*]] = tensor.subtensor_insert %[[sTD]] into %[[TC2]][{{.*}}] : tensor into tensor // CHECK: scf.yield %[[TD]] : tensor // CHECK: scf.yield %[[TD2]] : tensor // CHECK: scf.yield %[[TD1]] : tensor diff --git a/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir b/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir --- a/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir +++ b/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir @@ -16,11 +16,11 @@ %3 = scf.for %arg3 = %c0 to %0 step %c2 iter_args(%arg4 = %arg2) -> (tensor) { %4 = scf.for %arg5 = %c0 to %2 step %c3 iter_args(%arg6 = %arg4) -> (tensor) { %5 = scf.for %arg7 = %c0 to %1 step %c4 iter_args(%arg8 = %arg6) -> (tensor) { - %6 = subtensor %t0[%arg3, %arg7][%c2, 4][1, 1] : tensor to tensor - %7 = subtensor %arg1[%arg7, %arg5][4, %c3][1, 1] : tensor to tensor<4x?xf32> - %8 = subtensor %arg8[%arg3, %arg5][%c2, %c3][1, 1] : tensor to tensor + %6 = tensor.subtensor %t0[%arg3, %arg7][%c2, 4][1, 1] : tensor to tensor + %7 = tensor.subtensor %arg1[%arg7, %arg5][4, %c3][1, 1] : tensor to tensor<4x?xf32> + %8 = tensor.subtensor %arg8[%arg3, %arg5][%c2, %c3][1, 1] : tensor to tensor %9 = linalg.matmul ins(%6, %7 : tensor, tensor<4x?xf32>) outs(%8 : tensor) -> tensor - %10 = subtensor_insert %9 into %arg8[%arg3, %arg5] [%c2, %c3] [1, 1] : tensor into tensor + %10 = tensor.subtensor_insert %9 into %arg8[%arg3, %arg5] [%c2, %c3] [1, 1] : tensor into tensor scf.yield %10 : tensor } scf.yield %5 : tensor @@ -48,22 +48,22 @@ // CHECK-DAG: %[[dC1:.*]] = memref.dim %[[C]], %[[C1]] : tensor // CHECK: scf.for %[[I:[0-9a-z]*]] // CHECK: %[[sizeA0:.*]] = affine.min #[[BOUND2_MAP]](%[[I]])[%[[dA0]]] -// CHECK: %[[stA:.*]] = subtensor %[[A]][%[[I]], 0] [%[[sizeA0]], %[[dA1]]] [1, 1] : tensor to tensor +// CHECK: %[[stA:.*]] = tensor.subtensor %[[A]][%[[I]], 0] [%[[sizeA0]], %[[dA1]]] [1, 1] : tensor to tensor // CHECK: %[[sizeC0:.*]] = affine.min #[[BOUND2_MAP]](%[[I]])[%[[dC0]]] // CHECK-NEXT: scf.for %[[J:[0-9a-z]*]] // CHECK-NEXT: scf.for %[[K:[0-9a-z]*]] {{.*}} iter_args(%[[RES:[0-9a-z]*]] -// CHECK-DAG: %[[stB1:.*]] = subtensor %[[B]][%[[K]], %[[J]]] [4, 3] [1, 1] : tensor to tensor<4x3xf32> -// CHECK-DAG: %[[stF:.*]] = subtensor %[[RES]][%[[I]], %[[J]]] [2, 3] [1, 1] : tensor to tensor<2x3xf32> +// CHECK-DAG: %[[stB1:.*]] = tensor.subtensor %[[B]][%[[K]], %[[J]]] [4, 3] [1, 1] : tensor to tensor<4x3xf32> +// CHECK-DAG: %[[stF:.*]] = tensor.subtensor %[[RES]][%[[I]], %[[J]]] [2, 3] [1, 1] : tensor to tensor<2x3xf32> // // subtensors of the producing matmul. // CHECK: %[[sizeB1:.*]] = affine.min #[[BOUND4_MAP]](%[[K]])[%[[dB1]]] -// CHECK: %[[stB2:.*]] = subtensor %[[B]][0, %[[K]]] [%[[dB0]], %[[sizeB1]]] [1, 1] : tensor to tensor +// CHECK: %[[stB2:.*]] = tensor.subtensor %[[B]][0, %[[K]]] [%[[dB0]], %[[sizeB1]]] [1, 1] : tensor to tensor // CHECK: %[[sizeC1:.*]] = affine.min #[[BOUND4_MAP]](%[[K]])[%[[dC1]]] -// CHECK: %[[stC:.*]] = subtensor %[[C]][%[[I]], %[[K]]] [%[[sizeC0]], %[[sizeC1]]] [1, 1] : tensor to tensor +// CHECK: %[[stC:.*]] = tensor.subtensor %[[C]][%[[I]], %[[K]]] [%[[sizeC0]], %[[sizeC1]]] [1, 1] : tensor to tensor // CHECK: %[[stD:.*]] = linalg.matmul ins(%[[stA]], %[[stB2]] : tensor, tensor) outs(%[[stC]] : tensor) -> tensor // CHECK: %[[CAST:.*]] = tensor.cast %[[stD]] : tensor to tensor // CHECK-NEXT: %[[stG:.*]] = linalg.matmul ins(%[[CAST]], %[[stB1]] : tensor, tensor<4x3xf32>) outs(%[[stF]] : tensor<2x3xf32>) -> tensor<2x3xf32> -// CHECK-NEXT: subtensor_insert %[[stG]] into %[[RES]][%[[I]], %[[J]]] +// CHECK-NEXT: tensor.subtensor_insert %[[stG]] into %[[RES]][%[[I]], %[[J]]] // ----- @@ -87,9 +87,9 @@ %for0 = scf.for %iv0 = %c0 to %c112 step %c8 iter_args(%arg0 = %fill) -> tensor<1x112x112x32xf32> { %for1 = scf.for %iv1 = %c0 to %c112 step %c16 iter_args(%arg1 = %arg0) -> tensor<1x112x112x32xf32> { %for2 = scf.for %iv2 = %c0 to %c32 step %c4 iter_args(%arg2 = %arg1) -> tensor<1x112x112x32xf32> { - %0 = subtensor %conv[0, %iv0, %iv1, %iv2][1, 8, 16, 4][1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32> - %1 = subtensor %elementwise[0, %iv0, %iv1, %iv2][1, 8, 16, 4][1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32> - %2 = subtensor %arg2[0, %iv0, %iv1, %iv2][1, 8, 16, 4][1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32> + %0 = tensor.subtensor %conv[0, %iv0, %iv1, %iv2][1, 8, 16, 4][1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32> + %1 = tensor.subtensor %elementwise[0, %iv0, %iv1, %iv2][1, 8, 16, 4][1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32> + %2 = tensor.subtensor %arg2[0, %iv0, %iv1, %iv2][1, 8, 16, 4][1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32> %add = linalg.generic { indexing_maps = [ @@ -104,7 +104,7 @@ linalg.yield %result : f32 } -> tensor<1x8x16x4xf32> - %insert = subtensor_insert %add into %arg2[0, %iv0, %iv1, %iv2] [1, 8, 16, 4] [1, 1, 1, 1] : tensor<1x8x16x4xf32> into tensor<1x112x112x32xf32> + %insert = tensor.subtensor_insert %add into %arg2[0, %iv0, %iv1, %iv2] [1, 8, 16, 4] [1, 1, 1, 1] : tensor<1x8x16x4xf32> into tensor<1x112x112x32xf32> scf.yield %insert : tensor<1x112x112x32xf32> } scf.yield %for2 : tensor<1x112x112x32xf32> @@ -127,19 +127,19 @@ // CHECK-NEXT: %[[OFFSET_H:.+]] = affine.apply #[[MAP0]](%[[IV0]]) // CHECK-NEXT: scf.for %[[IV1:.+]] = %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[ARG1:.+]] = %[[ARG0]]) // CHECK-NEXT: %[[OFFSET_W:.+]] = affine.apply #[[MAP0]](%[[IV1]]) -// CHECK-NEXT: %[[ST_INPUT:.+]] = subtensor %arg0[0, %[[OFFSET_H]], %[[OFFSET_W]], 0] [1, 17, 33, 3] [1, 1, 1, 1] : tensor<1x225x225x3xf32> to tensor<1x17x33x3xf32> +// CHECK-NEXT: %[[ST_INPUT:.+]] = tensor.subtensor %arg0[0, %[[OFFSET_H]], %[[OFFSET_W]], 0] [1, 17, 33, 3] [1, 1, 1, 1] : tensor<1x225x225x3xf32> to tensor<1x17x33x3xf32> // CHECK-NEXT: scf.for %[[IV2:.+]] = %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[ARG2:.+]] = %[[ARG1]]) -// CHECK-NEXT: %[[ST_ELEM:.+]] = subtensor %[[ELEM]][0, %[[IV0]], %[[IV1]], %[[IV2]]] [1, 8, 16, 4] [1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32> -// CHECK-NEXT: %[[ST_ARG2:.+]] = subtensor %[[ARG2]][0, %[[IV0]], %[[IV1]], %[[IV2]]] [1, 8, 16, 4] [1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32> -// CHECK-NEXT: %[[ST_FILTER:.+]] = subtensor %[[FILTER]][0, 0, 0, %[[IV2]]] [3, 3, 3, 4] [1, 1, 1, 1] : tensor<3x3x3x32xf32> to tensor<3x3x3x4xf32> -// CHECK-NEXT: %[[ST_FILL:.+]] = subtensor %[[FILL]][0, %[[IV0]], %[[IV1]], %[[IV2]]] [1, 8, 16, 4] [1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32> +// CHECK-NEXT: %[[ST_ELEM:.+]] = tensor.subtensor %[[ELEM]][0, %[[IV0]], %[[IV1]], %[[IV2]]] [1, 8, 16, 4] [1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32> +// CHECK-NEXT: %[[ST_ARG2:.+]] = tensor.subtensor %[[ARG2]][0, %[[IV0]], %[[IV1]], %[[IV2]]] [1, 8, 16, 4] [1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32> +// CHECK-NEXT: %[[ST_FILTER:.+]] = tensor.subtensor %[[FILTER]][0, 0, 0, %[[IV2]]] [3, 3, 3, 4] [1, 1, 1, 1] : tensor<3x3x3x32xf32> to tensor<3x3x3x4xf32> +// CHECK-NEXT: %[[ST_FILL:.+]] = tensor.subtensor %[[FILL]][0, %[[IV0]], %[[IV1]], %[[IV2]]] [1, 8, 16, 4] [1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32> // CHECK-NEXT: %[[ST_CONV:.+]] = linalg.conv_2d_input_nhwc_filter_hwcf // CHECK-SAME: ins(%[[ST_INPUT]], %[[ST_FILTER]] : tensor<1x17x33x3xf32>, tensor<3x3x3x4xf32>) // CHECK-SAME: outs(%[[ST_FILL]] : tensor<1x8x16x4xf32>) // CHECK-NEXT: %[[ADD:.+]] = linalg.generic // CHECK-SAME: ins(%[[ST_CONV]], %[[ST_ELEM]] : tensor<1x8x16x4xf32>, tensor<1x8x16x4xf32>) // CHECK-SAME: outs(%[[ST_ARG2]] : tensor<1x8x16x4xf32>) -// CHECK: subtensor_insert %[[ADD]] into %[[ARG2]][0, %[[IV0]], %[[IV1]], %[[IV2]]] [1, 8, 16, 4] +// CHECK: tensor.subtensor_insert %[[ADD]] into %[[ARG2]][0, %[[IV0]], %[[IV1]], %[[IV2]]] [1, 8, 16, 4] // ----- @@ -174,9 +174,9 @@ %oh_size = affine.min affine_map<(d0)[s0] -> (16, -d0 + s0)>(%iv1)[%oh] %ow_size = affine.min affine_map<(d0)[s0] -> (4, -d0 + s0)>(%iv2)[%ow] %oc_size = affine.min affine_map<(d0)[s0] -> (2, -d0 + s0)>(%iv2)[%oc] - %0 = subtensor %conv[%iv0, %iv1, %iv2, %iv3][%n_size, %oh_size, %ow_size, %oc_size][1, 1, 1, 1] : tensor to tensor - %1 = subtensor %elementwise[%iv0, %iv1, %iv2, %iv3][%n_size, %oh_size, %ow_size, %oc_size][1, 1, 1, 1] : tensor to tensor - %2 = subtensor %arg3[%iv0, %iv1, %iv2, %iv3][%n_size, %oh_size, %ow_size, %oc_size][1, 1, 1, 1] : tensor to tensor + %0 = tensor.subtensor %conv[%iv0, %iv1, %iv2, %iv3][%n_size, %oh_size, %ow_size, %oc_size][1, 1, 1, 1] : tensor to tensor + %1 = tensor.subtensor %elementwise[%iv0, %iv1, %iv2, %iv3][%n_size, %oh_size, %ow_size, %oc_size][1, 1, 1, 1] : tensor to tensor + %2 = tensor.subtensor %arg3[%iv0, %iv1, %iv2, %iv3][%n_size, %oh_size, %ow_size, %oc_size][1, 1, 1, 1] : tensor to tensor %add = linalg.generic { indexing_maps = [ @@ -191,7 +191,7 @@ linalg.yield %result : f32 } -> tensor - %insert = subtensor_insert %add into %arg3[%iv0, %iv1, %iv2, %iv3] [%n_size, %oh_size, %ow_size, %oc_size] [1, 1, 1, 1] : tensor into tensor + %insert = tensor.subtensor_insert %add into %arg3[%iv0, %iv1, %iv2, %iv3] [%n_size, %oh_size, %ow_size, %oc_size] [1, 1, 1, 1] : tensor into tensor scf.yield %insert : tensor } scf.yield %for3 : tensor @@ -257,19 +257,19 @@ // CHECK-NEXT: %[[SIZE_ELEM_OC:.+]] = affine.min #[[BOUND2_MAP]](%[[IV2]])[%[[ELEM_OC]]] // CHECK-NEXT: %[[OFFSET_OW:.+]] = affine.apply #[[X2_MAP]](%[[IV2]]) // CHECK-NEXT: %[[SIZE_INPUT_W:.+]] = affine.min #[[INPUT_BOUND]](%[[SIZE_ELEM_OW]], %[[IV2]])[%[[FILTER_W]], %[[INPUT_W]]] -// CHECK-NEXT: %[[ST_INPUT:.+]] = subtensor %[[INPUT]][%[[IV0]], %[[OFFSET_OH]], %[[OFFSET_OW]], 0] +// CHECK-NEXT: %[[ST_INPUT:.+]] = tensor.subtensor %[[INPUT]][%[[IV0]], %[[OFFSET_OH]], %[[OFFSET_OW]], 0] // CHECK-SAME: [%[[SIZE_INPUT_N]], %[[SIZE_INPUT_H]], %[[SIZE_INPUT_W]], %[[INPUT_C]]] // CHECK-NEXT: %[[SIZE_ELEM_OW_2:.+]] = affine.min #[[BOUND4_MAP_2]](%[[IV2]])[%[[FILL_W]], %[[ELEM_OW]]] // CHECK-NEXT: scf.for %[[IV3:.+]] = %{{.+}} to %[[ELEM_OC]] step %{{.+}} iter_args(%[[ARG:[a-z0-9]+]] -// CHECK-NEXT: %[[ST_ELEM:.+]] = subtensor %[[ELEM]][%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]] +// CHECK-NEXT: %[[ST_ELEM:.+]] = tensor.subtensor %[[ELEM]][%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]] // CHECK-SAME: [%[[SIZE_ELEM_N]], %[[SIZE_ELEM_OH]], %[[SIZE_ELEM_OW]], %[[SIZE_ELEM_OC]]] -// CHECK-NEXT: %[[ST_ARG:.+]] = subtensor %[[ARG]][%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]] +// CHECK-NEXT: %[[ST_ARG:.+]] = tensor.subtensor %[[ARG]][%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]] // CHECK-SAME: [%[[SIZE_ELEM_N]], %[[SIZE_ELEM_OH]], %[[SIZE_ELEM_OW]], %[[SIZE_ELEM_OC]]] // CHECK-NEXT: %[[SIZE_ELEM_OC_2:.+]] = affine.min #[[BOUND2_MAP_2]](%[[IV3]], %[[IV2]])[%[[FILTER_OC]], %[[ELEM_OC]]] -// CHECK-NEXT: %[[ST_FILTER:.+]] = subtensor %[[FILTER]][0, 0, 0, %[[IV3]]] +// CHECK-NEXT: %[[ST_FILTER:.+]] = tensor.subtensor %[[FILTER]][0, 0, 0, %[[IV3]]] // CHECK-SAME: [%[[FILTER_H]], %[[FILTER_W]], %[[FILTER_IC]], %[[SIZE_ELEM_OC_2]]] // CHECK-NEXT: %[[SIZE_ELEM_OC_3:.+]] = affine.min #[[BOUND2_MAP_2]](%[[IV3]], %[[IV2]])[%[[FILL_C]], %[[ELEM_OC]]] -// CHECK-NEXT: %[[ST_FILL:.+]] = subtensor %[[FILL]][%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]] +// CHECK-NEXT: %[[ST_FILL:.+]] = tensor.subtensor %[[FILL]][%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]] // CHECK-SAME: [%[[SIZE_ELEM_N_2]], %[[SIZE_ELEM_OH_2]], %[[SIZE_ELEM_OW_2]], %[[SIZE_ELEM_OC_3]]] // CHECK-NEXT: %[[ST_CONV:.+]] = linalg.conv_2d_input_nhwc_filter_hwcf // CHECK-SAME: ins(%[[ST_INPUT]], %[[ST_FILTER]] : tensor, tensor) @@ -277,5 +277,5 @@ // CHECK-NEXT: %[[ST_ADD:.+]] = linalg.generic // CHECK-SAME: ins(%[[ST_CONV]], %[[ST_ELEM]] : tensor, tensor) // CHECK-SAME: outs(%[[ST_ARG]] : tensor) -// CHECK: subtensor_insert %[[ST_ADD]] into %[[ARG]][%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]] +// CHECK: tensor.subtensor_insert %[[ST_ADD]] into %[[ARG]][%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]] // CHECK-SAME: [%[[SIZE_ELEM_N]], %[[SIZE_ELEM_OH]], %[[SIZE_ELEM_OW]], %[[SIZE_ELEM_OC]]] diff --git a/mlir/test/Dialect/Linalg/tile-and-pad-tensors.mlir b/mlir/test/Dialect/Linalg/tile-and-pad-tensors.mlir --- a/mlir/test/Dialect/Linalg/tile-and-pad-tensors.mlir +++ b/mlir/test/Dialect/Linalg/tile-and-pad-tensors.mlir @@ -12,9 +12,9 @@ // CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor) { // CHECK: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor) { // CHECK: %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor) { -// CHECK: %[[sTA:.*]] = subtensor %[[TA]][{{.*}}] : tensor to tensor -// CHECK: %[[sTB:.*]] = subtensor %[[TB]][{{.*}}] : tensor to tensor -// CHECK: %[[sTC:.*]] = subtensor %[[TC2]][{{.*}}] : tensor to tensor +// CHECK: %[[sTA:.*]] = tensor.subtensor %[[TA]][{{.*}}] : tensor to tensor +// CHECK: %[[sTB:.*]] = tensor.subtensor %[[TB]][{{.*}}] : tensor to tensor +// CHECK: %[[sTC:.*]] = tensor.subtensor %[[TC2]][{{.*}}] : tensor to tensor // Dynamic op has been canonicalized away. // CHECK-NOT: linalg.matmul {{.*}} tensor @@ -28,8 +28,8 @@ // CHECK: : tensor to tensor<2x3xi32> // CHECK: %[[pD:.*]] = linalg.matmul_i8_i8_i32 ins(%[[pA]], %[[pB]] : tensor<2x4xi8>, tensor<4x3xi8>) // CHECK-SAME: outs(%[[pC]] : tensor<2x3xi32>) -> tensor<2x3xi32> -// CHECK: %[[sTD:.*]] = subtensor %[[pD]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : tensor<2x3xi32> to tensor -// CHECK: %[[TD:.*]] = subtensor_insert %[[sTD]] into %[[TC2]][{{.*}}] : tensor into tensor +// CHECK: %[[sTD:.*]] = tensor.subtensor %[[pD]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : tensor<2x3xi32> to tensor +// CHECK: %[[TD:.*]] = tensor.subtensor_insert %[[sTD]] into %[[TC2]][{{.*}}] : tensor into tensor // CHECK: scf.yield %[[TD]] : tensor // CHECK: scf.yield %[[TD2]] : tensor // CHECK: scf.yield %[[TD1]] : tensor @@ -52,15 +52,15 @@ // CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor) { // CHECK: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor) { // CHECK: %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor) { -// CHECK: %[[sTC:.*]] = subtensor %[[TC2]][{{.*}}] : tensor to tensor +// CHECK: %[[sTC:.*]] = tensor.subtensor %[[TC2]][{{.*}}] : tensor to tensor // Padding injects static information. // CHECK: %[[pC:.*]] = linalg.pad_tensor %[[sTC]] low[%[[C0]], %[[C0]], %[[C0]]] high[%{{.*}}, %{{.*}}, %{{.*}}] // CHECK: : tensor to tensor<2x3x4xf32> // CHECK: %[[pD:.*]] = linalg.generic // CHECK-SAME: ins(%[[VAL]] : f32) outs(%[[pC]] : tensor<2x3x4xf32>) -// CHECK: %[[sTD:.*]] = subtensor %[[pD]][0, 0, 0] [%{{.*}}, %{{.*}}, %{{.*}}] [1, 1, 1] : tensor<2x3x4xf32> to tensor -// CHECK: %[[TD:.*]] = subtensor_insert %[[sTD]] into %[[TC2]][{{.*}}] : tensor into tensor +// CHECK: %[[sTD:.*]] = tensor.subtensor %[[pD]][0, 0, 0] [%{{.*}}, %{{.*}}, %{{.*}}] [1, 1, 1] : tensor<2x3x4xf32> to tensor +// CHECK: %[[TD:.*]] = tensor.subtensor_insert %[[sTD]] into %[[TC2]][{{.*}}] : tensor into tensor // CHECK: scf.yield %[[TD]] : tensor // CHECK: scf.yield %[[TD2]] : tensor // CHECK: scf.yield %[[TD1]] : tensor @@ -104,11 +104,11 @@ // CHECK-1DIM-TILE: %[[C0:.*]] = constant 0 : index // CHECK-1DIM-TILE: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor) { // CHECK-1DIM-TILE: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor) { -// CHECK-1DIM-TILE: %[[sTA:.*]] = subtensor %[[TA]][{{.*}}] : tensor to tensor +// CHECK-1DIM-TILE: %[[sTA:.*]] = tensor.subtensor %[[TA]][{{.*}}] : tensor to tensor // CHECK-1DIM-TILE: %[[sTAc:.*]] = tensor.cast %[[sTA]] : tensor to tensor -// CHECK-1DIM-TILE: %[[sTB:.*]] = subtensor %[[TB]][{{.*}}] : tensor<8x?xi8> to tensor<8x?xi8> +// CHECK-1DIM-TILE: %[[sTB:.*]] = tensor.subtensor %[[TB]][{{.*}}] : tensor<8x?xi8> to tensor<8x?xi8> // CHECK-1DIM-TILE: %[[sTBc:.*]] = tensor.cast %[[sTB]] : tensor<8x?xi8> to tensor -// CHECK-1DIM-TILE: %[[sTC:.*]] = subtensor %[[TC1]][{{.*}}] : tensor to tensor +// CHECK-1DIM-TILE: %[[sTC:.*]] = tensor.subtensor %[[TC1]][{{.*}}] : tensor to tensor // CHECK-1DIM-TILE: %[[pA:.*]] = linalg.pad_tensor %[[sTAc]] low[%[[C0]], %[[C0]]] high[%{{.*}}, %{{.*}}] // CHECK-1DIM-TILE: : tensor to tensor<2x8xi8> // CHECK-1DIM-TILE: %[[pB:.*]] = linalg.pad_tensor %[[sTBc]] low[%[[C0]], %[[C0]]] high[%{{.*}}, %{{.*}}] diff --git a/mlir/test/Dialect/Linalg/tile-tensors.mlir b/mlir/test/Dialect/Linalg/tile-tensors.mlir --- a/mlir/test/Dialect/Linalg/tile-tensors.mlir +++ b/mlir/test/Dialect/Linalg/tile-tensors.mlir @@ -11,12 +11,12 @@ // CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor) { // CHECK: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor) { // CHECK: %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor) { -// CHECK: %[[sTA:.*]] = subtensor %[[TA]][{{.*}}] : tensor to tensor -// CHECK: %[[sTB:.*]] = subtensor %[[TB]][{{.*}}] : tensor to tensor -// CHECK: %[[sTC:.*]] = subtensor %[[TC2]][{{.*}}] : tensor to tensor +// CHECK: %[[sTA:.*]] = tensor.subtensor %[[TA]][{{.*}}] : tensor to tensor +// CHECK: %[[sTB:.*]] = tensor.subtensor %[[TB]][{{.*}}] : tensor to tensor +// CHECK: %[[sTC:.*]] = tensor.subtensor %[[TC2]][{{.*}}] : tensor to tensor // CHECK: %[[sTD:.*]] = linalg.matmul ins(%[[sTA]], %[[sTB]] : tensor, tensor) // CHECK-SAME: outs(%[[sTC]] : tensor) -> tensor -// CHECK: %[[TD:.*]] = subtensor_insert %[[sTD]] into %[[TC2]][{{.*}}] : tensor into tensor +// CHECK: %[[TD:.*]] = tensor.subtensor_insert %[[sTD]] into %[[TC2]][{{.*}}] : tensor into tensor // CHECK: scf.yield %[[TD]] : tensor // CHECK: scf.yield %[[TD2]] : tensor // CHECK: scf.yield %[[TD1]] : tensor @@ -51,14 +51,14 @@ // TLOOP-SAME: iterators["parallel", "parallel", "reduction"] // TLOOP-SAME: distribution["block_x", "block_y", "none"] { -// TLOOP: %[[SUB_ARG_0:.*]] = subtensor %[[A0]][%[[I]], %[[K]]] -// TLOOP: %[[SUB_ARG_1:.*]] = subtensor %[[A1]][%[[K]], %[[J]]] -// TLOOP: %[[SUB_ARG_2:.*]] = subtensor %[[A2]][%[[I]], %[[J]]] +// TLOOP: %[[SUB_ARG_0:.*]] = tensor.subtensor %[[A0]][%[[I]], %[[K]]] +// TLOOP: %[[SUB_ARG_1:.*]] = tensor.subtensor %[[A1]][%[[K]], %[[J]]] +// TLOOP: %[[SUB_ARG_2:.*]] = tensor.subtensor %[[A2]][%[[I]], %[[J]]] // TLOOP: %[[PROD:.*]] = linalg.matmul ins(%[[SUB_ARG_0]], %[[SUB_ARG_1]] // TLOOP-SE: outs(%[[SUB_ARG_2]] : [[TY]]) -> [[TY]] -// TLOOP: %[[O:.*]] = subtensor_insert %[[PROD]] into %[[A2]][%[[I]], %[[J]]] +// TLOOP: %[[O:.*]] = tensor.subtensor_insert %[[PROD]] into %[[A2]][%[[I]], %[[J]]] // TLOOP: linalg.yield %[[O]] : [[TY]] // ----- @@ -93,13 +93,13 @@ // CHECK: %[[TD0:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC0:.+]] = %[[INIT]]) -> (tensor) { // CHECK: %[[TD1:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC1:.+]] = %[[TC0]]) -> (tensor) { // CHECK: %[[TD2:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC2:.+]] = %[[TC1]]) -> (tensor) { -// CHECK: %[[STARG0:.+]] = subtensor %[[ARG0]][{{.+}}] : tensor to tensor -// CHECK: %[[STARG1:.+]] = subtensor %[[ARG1]][{{.+}}] : tensor to tensor -// CHECK: %[[STARG2:.+]] = subtensor %[[TC2]][{{.+}}] : tensor to tensor +// CHECK: %[[STARG0:.+]] = tensor.subtensor %[[ARG0]][{{.+}}] : tensor to tensor +// CHECK: %[[STARG1:.+]] = tensor.subtensor %[[ARG1]][{{.+}}] : tensor to tensor +// CHECK: %[[STARG2:.+]] = tensor.subtensor %[[TC2]][{{.+}}] : tensor to tensor // CHECK: %[[STRETURN:.+]] = linalg.generic // CHECK-SAME: ins(%[[STARG0]], %[[STARG1]] : tensor, tensor) // CHECK-SAME: outs(%[[STARG2]] : tensor) -// CHECK: %[[TD:.+]] = subtensor_insert %[[STRETURN]] into %[[TC2]] +// CHECK: %[[TD:.+]] = tensor.subtensor_insert %[[STRETURN]] into %[[TC2]] // CHECK: scf.yield %[[TD]] // CHECK: } // CHECK: scf.yield %[[TD2]] diff --git a/mlir/test/Dialect/Linalg/vectorization.mlir b/mlir/test/Dialect/Linalg/vectorization.mlir --- a/mlir/test/Dialect/Linalg/vectorization.mlir +++ b/mlir/test/Dialect/Linalg/vectorization.mlir @@ -586,7 +586,7 @@ // CHECK: %[[INIT:.*]] = linalg.init_tensor [6, %[[V1]], %[[V2]], %[[V5]]] : tensor<6x?x?x?xf32> // CHECK: %[[FILL:.*]] = linalg.fill(%[[INIT]], %{{.*}}) : tensor<6x?x?x?xf32>, f32 -> tensor<6x?x?x?xf32> // CHECK: %[[SRCDIM:.*]] = memref.dim %[[SRC]], %[[C3]] : tensor<1x2x2x?xf32> -// CHECK: %[[RESULT:.*]] = subtensor_insert %[[SRC]] into %[[FILL]][2, %[[LOW]], 3, 3] [1, 2, 2, %[[SRCDIM]]] [1, 1, 1, 1] : tensor<1x2x2x?xf32> into tensor<6x?x?x?xf32> +// CHECK: %[[RESULT:.*]] = tensor.subtensor_insert %[[SRC]] into %[[FILL]][2, %[[LOW]], 3, 3] [1, 2, 2, %[[SRCDIM]]] [1, 1, 1, 1] : tensor<1x2x2x?xf32> into tensor<6x?x?x?xf32> // CHECK: return %[[RESULT]] func @pad_static_dynamic(%arg0: tensor<1x2x2x?xf32>, %low: index, %high: index, %pad_value: f32) -> tensor<6x?x?x?xf32> { @@ -638,7 +638,7 @@ } : tensor<5x6xf32> to tensor<10x13xf32> %1 = vector.transfer_write %arg1, %0[%c0, %c0] : vector<7x9xf32>, tensor<10x13xf32> - %2 = subtensor %1[0, 0] [5, 6] [1, 1] : tensor<10x13xf32> to tensor<5x6xf32> + %2 = tensor.subtensor %1[0, 0] [5, 6] [1, 1] : tensor<10x13xf32> to tensor<5x6xf32> return %2 : tensor<5x6xf32> } @@ -648,14 +648,14 @@ // CHECK-SAME: %[[ARG0:.*]]: tensor, %[[ARG1:.*]]: vector<7x9xf32>, %[[SIZE:.*]]: index, %[[PADDING:.*]]: index // CHECK-NOT: linalg.pad_tensor // CHECK: %[[C0:.*]] = constant 0 : index -// CHECK: %[[SUB:.*]] = subtensor %[[ARG0]][0, 0] [%[[SIZE]], 6] [1, 1] : tensor to tensor +// CHECK: %[[SUB:.*]] = tensor.subtensor %[[ARG0]][0, 0] [%[[SIZE]], 6] [1, 1] : tensor to tensor // CHECK: %[[RESULT:.*]] = vector.transfer_write %[[ARG1]], %[[SUB]][%[[C0]], %[[C0]]] : vector<7x9xf32>, tensor // CHECK: return %[[RESULT]] func @pad_and_transfer_write_dynamic_static( %arg0: tensor, %arg1: vector<7x9xf32>, %size: index, %padding: index) -> tensor { %c0 = constant 0 : index %c5 = constant 5.0 : f32 - %s = subtensor %arg0[0, 0] [%size, 6] [1, 1] + %s = tensor.subtensor %arg0[0, 0] [%size, 6] [1, 1] : tensor to tensor %0 = linalg.pad_tensor %s low[0, 0] high[%padding, 7] { ^bb0(%arg2: index, %arg3: index): @@ -663,7 +663,7 @@ } : tensor to tensor %1 = vector.transfer_write %arg1, %0[%c0, %c0] : vector<7x9xf32>, tensor - %2 = subtensor %1[0, 0] [%size, 6] [1, 1] : tensor to tensor + %2 = tensor.subtensor %1[0, 0] [%size, 6] [1, 1] : tensor to tensor return %2 : tensor } @@ -685,7 +685,7 @@ ^bb0(%arg2: index, %arg3: index): linalg.yield %c5 : f32 } : tensor<5x6xf32> to tensor<7x9xf32> - %r = subtensor_insert %0 into %arg1[0, 0][7, 9][1, 1] : tensor<7x9xf32> into tensor<12x13xf32> + %r = tensor.subtensor_insert %0 into %arg1[0, 0][7, 9][1, 1] : tensor<7x9xf32> into tensor<12x13xf32> return %r : tensor<12x13xf32> } diff --git a/mlir/test/Dialect/MemRef/canonicalize.mlir b/mlir/test/Dialect/MemRef/canonicalize.mlir --- a/mlir/test/Dialect/MemRef/canonicalize.mlir +++ b/mlir/test/Dialect/MemRef/canonicalize.mlir @@ -217,3 +217,153 @@ return } +// ----- + +// Test case: Folding of memref.load(memref.buffer_cast(%v, %idxs)) +// -> tensor.extract(%v, %idx) +// CHECK-LABEL: func @load_from_buffer_cast( +// CHECK-SAME: %[[IDX0:[0-9a-z]+]]: index, %[[IDX1:[0-9a-z]+]]: index +// CHECK-SAME: %[[TENSOR:[0-9a-z]+]]: tensor +// CHECK: %[[RES:.*]] = tensor.extract %[[TENSOR]][%[[IDX0]], %[[IDX1]]] +// CHECK-NOT: memref.load +// CHECK: return %[[RES]] : f32 +func @load_from_buffer_cast(%arg0: index, %arg1: index, %arg2: tensor) -> f32 { + %0 = memref.buffer_cast %arg2 : memref + %1 = memref.load %0[%arg0, %arg1] : memref + return %1 : f32 +} + +// ----- + + +// Test case: Basic folding of memref.dim(memref.tensor_load(m)) -> memref.dim(m). +// CHECK-LABEL: func @dim_of_tensor_load( +// CHECK-SAME: %[[MEMREF:[0-9a-z]*]]: memref +// CHECK: %[[C0:.*]] = constant 0 +// CHECK: %[[D:.*]] = memref.dim %[[MEMREF]], %[[C0]] +// CHECK: return %[[D]] : index +func @dim_of_tensor_load(%arg0: memref) -> index { + %c0 = constant 0 : index + %0 = memref.tensor_load %arg0 : memref + %1 = memref.dim %0, %c0 : tensor + return %1 : index +} + +// ----- + +// Test case: Folding of memref.dim(tensor.generate %idx) -> %idx +// CHECK-LABEL: func @dim_of_tensor.generate( +// CHECK-SAME: %[[IDX0:[0-9a-z]+]]: index, %[[IDX1:[0-9a-z]+]]: index +// CHECK-NOT: memref.dim +// CHECK: return %[[IDX1]] : index +func @dim_of_tensor.generate(%arg0: index, %arg1: index) -> index { + %c3 = constant 3 : index + %0 = tensor.generate %arg0, %arg1 { + ^bb0(%arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index): + tensor.yield %c3 : index + } : tensor<2x?x4x?x5xindex> + %1 = memref.dim %0, %c3 : tensor<2x?x4x?x5xindex> + return %1 : index +} + +// ----- + +// Test case: Folding of memref.dim(memref.alloca(%size), %idx) -> %size +// CHECK-LABEL: func @dim_of_alloca( +// CHECK-SAME: %[[SIZE:[0-9a-z]+]]: index +// CHECK-NEXT: return %[[SIZE]] : index +func @dim_of_alloca(%size: index) -> index { + %0 = memref.alloca(%size) : memref + %c0 = constant 0 : index + %1 = memref.dim %0, %c0 : memref + return %1 : index +} + +// ----- + +// Test case: Folding of memref.dim(memref.alloca(rank(%v)), %idx) -> rank(%v) +// CHECK-LABEL: func @dim_of_alloca_with_dynamic_size( +// CHECK-SAME: %[[MEM:[0-9a-z]+]]: memref<*xf32> +// CHECK-NEXT: %[[RANK:.*]] = rank %[[MEM]] : memref<*xf32> +// CHECK-NEXT: return %[[RANK]] : index +func @dim_of_alloca_with_dynamic_size(%arg0: memref<*xf32>) -> index { + %0 = rank %arg0 : memref<*xf32> + %1 = memref.alloca(%0) : memref + %c0 = constant 0 : index + %2 = memref.dim %1, %c0 : memref + return %2 : index +} + +// ----- + +// Test case: Folding of memref.dim(memref.reshape %v %shp, %idx) -> memref.load %shp[%idx] +// CHECK-LABEL: func @dim_of_memref_reshape( +// CHECK-SAME: %[[MEM:[0-9a-z]+]]: memref<*xf32>, +// CHECK-SAME: %[[SHP:[0-9a-z]+]]: memref +// CHECK-NEXT: %[[IDX:.*]] = constant 3 +// CHECK-NEXT: %[[DIM:.*]] = memref.load %[[SHP]][%[[IDX]]] +// CHECK-NEXT: memref.store +// CHECK-NOT: memref.dim +// CHECK: return %[[DIM]] : index +func @dim_of_memref_reshape(%arg0: memref<*xf32>, %arg1: memref) + -> index { + %c3 = constant 3 : index + %0 = memref.reshape %arg0(%arg1) + : (memref<*xf32>, memref) -> memref<*xf32> + // Update the shape to test that he load ends up in the right place. + memref.store %c3, %arg1[%c3] : memref + %1 = memref.dim %0, %c3 : memref<*xf32> + return %1 : index +} + +// ----- + +// Test case: Folding of memref.dim(memref.reshape %v %shp, %idx) -> memref.load %shp[%idx] +// CHECK-LABEL: func @dim_of_memref_reshape_i32( +// CHECK-SAME: %[[MEM:[0-9a-z]+]]: memref<*xf32>, +// CHECK-SAME: %[[SHP:[0-9a-z]+]]: memref +// CHECK-NEXT: %[[IDX:.*]] = constant 3 +// CHECK-NEXT: %[[DIM:.*]] = memref.load %[[SHP]][%[[IDX]]] +// CHECK-NEXT: %[[CAST:.*]] = index_cast %[[DIM]] +// CHECK-NOT: memref.dim +// CHECK: return %[[CAST]] : index +func @dim_of_memref_reshape_i32(%arg0: memref<*xf32>, %arg1: memref) + -> index { + %c3 = constant 3 : index + %0 = memref.reshape %arg0(%arg1) + : (memref<*xf32>, memref) -> memref<*xf32> + %1 = memref.dim %0, %c3 : memref<*xf32> + return %1 : index +} + +// ----- + +// Test case: Folding memref.dim(tensor.cast %0, %idx) -> memref.dim %0, %idx +// CHECK-LABEL: func @fold_dim_of_tensor.cast +// CHECK-SAME: %[[ARG0:.[a-z0-9A-Z_]+]]: tensor<4x?xf32> +// CHECK-DAG: %[[C1:.+]] = constant 1 : index +// CHECK-DAG: %[[C4:.+]] = constant 4 : index +// CHECK: %[[T0:.+]] = memref.dim %[[ARG0]], %[[C1]] +// CHECK-NEXT: return %[[C4]], %[[T0]] +func @fold_dim_of_tensor.cast(%arg0 : tensor<4x?xf32>) -> (index, index) { + %c0 = constant 0 : index + %c1 = constant 1 : index + %0 = tensor.cast %arg0 : tensor<4x?xf32> to tensor + %1 = memref.dim %0, %c0 : tensor + %2 = memref.dim %0, %c1 : tensor + return %1, %2: index, index +} + +// ----- + +// CHECK-LABEL: func @tensor_cast_to_memref +// CHECK-SAME: %[[ARG0:.+]]: tensor<4x6x16x32xi8> +// CHECK: %[[M:.+]] = memref.buffer_cast %[[ARG0]] : memref<4x6x16x32xi8> +// CHECK: %[[M1:.+]] = memref.cast %[[M]] : memref<4x6x16x32xi8> to memref +// CHECK: return %[[M1]] : memref +func @tensor_cast_to_memref(%arg0 : tensor<4x6x16x32xi8>) -> + memref { + %0 = tensor.cast %arg0 : tensor<4x6x16x32xi8> to tensor + %1 = memref.buffer_cast %0 : memref + return %1 : memref +} diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir --- a/mlir/test/Dialect/SCF/canonicalize.mlir +++ b/mlir/test/Dialect/SCF/canonicalize.mlir @@ -659,10 +659,10 @@ scf.yield %2 : tensor } // CHECK-NOT: tensor.cast -// CHECK: %[[RES:.*]] = subtensor_insert %[[FOR_RES]] into %[[T1]][0, 0] [32, 1024] [1, 1] : tensor<32x1024xf32> into tensor<1024x1024xf32> +// CHECK: %[[RES:.*]] = tensor.subtensor_insert %[[FOR_RES]] into %[[T1]][0, 0] [32, 1024] [1, 1] : tensor<32x1024xf32> into tensor<1024x1024xf32> // CHECK: return %[[RES]] : tensor<1024x1024xf32> %2 = tensor.cast %1 : tensor to tensor<32x1024xf32> - %res = subtensor_insert %2 into %t1[0, 0] [32, 1024] [1, 1] : tensor<32x1024xf32> into tensor<1024x1024xf32> + %res = tensor.subtensor_insert %2 into %t1[0, 0] [32, 1024] [1, 1] : tensor<32x1024xf32> into tensor<1024x1024xf32> return %res : tensor<1024x1024xf32> } 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 @@ -1,53 +1,5 @@ // RUN: mlir-opt %s -canonicalize --split-input-file | FileCheck %s -// Test case: Basic folding of memref.dim(memref.tensor_load(m)) -> memref.dim(m). -// CHECK-LABEL: func @dim_of_tensor_load( -// CHECK-SAME: %[[MEMREF:[0-9a-z]*]]: memref -// CHECK: %[[C0:.*]] = constant 0 -// CHECK: %[[D:.*]] = memref.dim %[[MEMREF]], %[[C0]] -// CHECK: return %[[D]] : index -func @dim_of_tensor_load(%arg0: memref) -> index { - %c0 = constant 0 : index - %0 = memref.tensor_load %arg0 : memref - %1 = memref.dim %0, %c0 : tensor - return %1 : index -} - -// ----- - -// Test case: Folding of memref.load(memref.buffer_cast(%v, %idxs)) -// -> tensor.extract(%v, %idx) -// CHECK-LABEL: func @load_from_buffer_cast( -// CHECK-SAME: %[[IDX0:[0-9a-z]+]]: index, %[[IDX1:[0-9a-z]+]]: index -// CHECK-SAME: %[[TENSOR:[0-9a-z]+]]: tensor -// CHECK: %[[RES:.*]] = tensor.extract %[[TENSOR]][%[[IDX0]], %[[IDX1]]] -// CHECK-NOT: memref.load -// CHECK: return %[[RES]] : f32 -func @load_from_buffer_cast(%arg0: index, %arg1: index, %arg2: tensor) -> f32 { - %0 = memref.buffer_cast %arg2 : memref - %1 = memref.load %0[%arg0, %arg1] : memref - return %1 : f32 -} - -// ----- - -// Test case: Folding of memref.dim(tensor.generate %idx) -> %idx -// CHECK-LABEL: func @dim_of_tensor.generate( -// CHECK-SAME: %[[IDX0:[0-9a-z]+]]: index, %[[IDX1:[0-9a-z]+]]: index -// CHECK-NOT: memref.dim -// CHECK: return %[[IDX1]] : index -func @dim_of_tensor.generate(%arg0: index, %arg1: index) -> index { - %c3 = constant 3 : index - %0 = tensor.generate %arg0, %arg1 { - ^bb0(%arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index): - tensor.yield %c3 : index - } : tensor<2x?x4x?x5xindex> - %1 = memref.dim %0, %c3 : tensor<2x?x4x?x5xindex> - return %1 : index -} - -// ----- - // Test case: Folding of comparisons with equal operands. // CHECK-LABEL: @cmpi_equal_operands // CHECK-DAG: %[[T:.*]] = constant true @@ -72,327 +24,6 @@ // ----- -// Test case: Folding of memref.dim(memref.alloca(%size), %idx) -> %size -// CHECK-LABEL: func @dim_of_alloca( -// CHECK-SAME: %[[SIZE:[0-9a-z]+]]: index -// CHECK-NEXT: return %[[SIZE]] : index -func @dim_of_alloca(%size: index) -> index { - %0 = memref.alloca(%size) : memref - %c0 = constant 0 : index - %1 = memref.dim %0, %c0 : memref - return %1 : index -} - -// ----- - -// Test case: Folding of memref.dim(memref.alloca(rank(%v)), %idx) -> rank(%v) -// CHECK-LABEL: func @dim_of_alloca_with_dynamic_size( -// CHECK-SAME: %[[MEM:[0-9a-z]+]]: memref<*xf32> -// CHECK-NEXT: %[[RANK:.*]] = rank %[[MEM]] : memref<*xf32> -// CHECK-NEXT: return %[[RANK]] : index -func @dim_of_alloca_with_dynamic_size(%arg0: memref<*xf32>) -> index { - %0 = rank %arg0 : memref<*xf32> - %1 = memref.alloca(%0) : memref - %c0 = constant 0 : index - %2 = memref.dim %1, %c0 : memref - return %2 : index -} - -// ----- - -// Test case: Folding of memref.dim(memref.reshape %v %shp, %idx) -> memref.load %shp[%idx] -// CHECK-LABEL: func @dim_of_memref_reshape( -// CHECK-SAME: %[[MEM:[0-9a-z]+]]: memref<*xf32>, -// CHECK-SAME: %[[SHP:[0-9a-z]+]]: memref -// CHECK-NEXT: %[[IDX:.*]] = constant 3 -// CHECK-NEXT: %[[DIM:.*]] = memref.load %[[SHP]][%[[IDX]]] -// CHECK-NEXT: memref.store -// CHECK-NOT: memref.dim -// CHECK: return %[[DIM]] : index -func @dim_of_memref_reshape(%arg0: memref<*xf32>, %arg1: memref) - -> index { - %c3 = constant 3 : index - %0 = memref.reshape %arg0(%arg1) - : (memref<*xf32>, memref) -> memref<*xf32> - // Update the shape to test that he load ends up in the right place. - memref.store %c3, %arg1[%c3] : memref - %1 = memref.dim %0, %c3 : memref<*xf32> - return %1 : index -} - -// ----- - -// Test case: Folding of memref.dim(memref.reshape %v %shp, %idx) -> memref.load %shp[%idx] -// CHECK-LABEL: func @dim_of_memref_reshape_i32( -// CHECK-SAME: %[[MEM:[0-9a-z]+]]: memref<*xf32>, -// CHECK-SAME: %[[SHP:[0-9a-z]+]]: memref -// CHECK-NEXT: %[[IDX:.*]] = constant 3 -// CHECK-NEXT: %[[DIM:.*]] = memref.load %[[SHP]][%[[IDX]]] -// CHECK-NEXT: %[[CAST:.*]] = index_cast %[[DIM]] -// CHECK-NOT: memref.dim -// CHECK: return %[[CAST]] : index -func @dim_of_memref_reshape_i32(%arg0: memref<*xf32>, %arg1: memref) - -> index { - %c3 = constant 3 : index - %0 = memref.reshape %arg0(%arg1) - : (memref<*xf32>, memref) -> memref<*xf32> - %1 = memref.dim %0, %c3 : memref<*xf32> - return %1 : index -} - -// ----- - -// Test case: Folding memref.dim(tensor.cast %0, %idx) -> memref.dim %0, %idx -// CHECK-LABEL: func @fold_dim_of_tensor.cast -// CHECK-SAME: %[[ARG0:.[a-z0-9A-Z_]+]]: tensor<4x?xf32> -// CHECK-DAG: %[[C1:.+]] = constant 1 : index -// CHECK-DAG: %[[C4:.+]] = constant 4 : index -// CHECK: %[[T0:.+]] = memref.dim %[[ARG0]], %[[C1]] -// CHECK-NEXT: return %[[C4]], %[[T0]] -func @fold_dim_of_tensor.cast(%arg0 : tensor<4x?xf32>) -> (index, index) { - %c0 = constant 0 : index - %c1 = constant 1 : index - %0 = tensor.cast %arg0 : tensor<4x?xf32> to tensor - %1 = memref.dim %0, %c0 : tensor - %2 = memref.dim %0, %c1 : tensor - return %1, %2: index, index -} - -// ----- - -// CHECK-LABEL: func @tensor_cast_to_memref -// CHECK-SAME: %[[ARG0:.+]]: tensor<4x6x16x32xi8> -// CHECK: %[[M:.+]] = memref.buffer_cast %[[ARG0]] : memref<4x6x16x32xi8> -// CHECK: %[[M1:.+]] = memref.cast %[[M]] : memref<4x6x16x32xi8> to memref -// CHECK: return %[[M1]] : memref -func @tensor_cast_to_memref(%arg0 : tensor<4x6x16x32xi8>) -> - memref { - %0 = tensor.cast %arg0 : tensor<4x6x16x32xi8> to tensor - %1 = memref.buffer_cast %0 : memref - return %1 : memref -} - -// ----- - -func @subtensor_canonicalize(%arg0 : tensor, %arg1 : index, - %arg2 : index) -> tensor -{ - %c0 = constant 0 : index - %c1 = constant 1 : index - %c4 = constant 4 : index - %0 = subtensor %arg0[%c0, %arg1, %c1] [%c4, %c1, %arg2] [%c1, %c1, %c1] : tensor to tensor - return %0 : tensor -} -// CHECK-LABEL: func @subtensor_canonicalize -// CHECK-SAME: %[[ARG0:.+]]: tensor -// CHECK: %[[SUBTENSOR:.+]] = subtensor %[[ARG0]][0, %{{[a-zA-Z0-9_]+}}, 1] -// CHECK-SAME: [4, 1, %{{[a-zA-Z0-9_]+}}] [1, 1, 1] -// CHECK-SAME: : tensor to tensor<4x1x?xf32> -// CHECK: %[[RESULT:.+]] = tensor.cast %[[SUBTENSOR]] -// CHEKC: return %[[RESULT]] - -// ----- - -func @rank_reducing_subtensor_canonicalize(%arg0 : tensor, %arg1 : index, - %arg2 : index) -> tensor -{ - %c0 = constant 0 : index - %c1 = constant 1 : index - %c4 = constant 4 : index - %0 = subtensor %arg0[%c0, %arg1, %c1] [%c4, 1, %arg2] [%c1, %c1, %c1] : tensor to tensor - return %0 : tensor -} -// CHECK-LABEL: func @rank_reducing_subtensor_canonicalize -// CHECK-SAME: %[[ARG0:.+]]: tensor -// CHECK: %[[SUBTENSOR:.+]] = subtensor %[[ARG0]][0, %{{[a-zA-Z0-9_]+}}, 1] -// CHECK-SAME: [4, 1, %{{[a-zA-Z0-9_]+}}] [1, 1, 1] -// CHECK-SAME: : tensor to tensor<4x?xf32> -// CHECK: %[[RESULT:.+]] = tensor.cast %[[SUBTENSOR]] -// CHEKC: return %[[RESULT]] - -// ----- - -// CHECK-LABEL: func @trivial_subtensor -// CHECK-SAME: %[[ARG0:.[a-z0-9A-Z_]+]]: tensor<4x6x16x32xi8> -// CHECK-NOT: subtensor -// CHECK: return %[[ARG0]] : tensor<4x6x16x32xi8> -func @trivial_subtensor(%arg0 : tensor<4x6x16x32xi8>) -> tensor<4x6x16x32xi8> { - %0 = subtensor %arg0[0, 0, 0, 0] [4, 6, 16, 32] [1, 1, 1, 1] : tensor<4x6x16x32xi8> to tensor<4x6x16x32xi8> - return %0 : tensor<4x6x16x32xi8> -} - -// ----- - -// CHECK-LABEL: func @trivial_subtensor_insert -// CHECK-SAME: %[[ARG0:.[a-z0-9A-Z_]+]]: tensor<4x6x16x32xi8> -// CHECK-NOT: subtensor -// CHECK: return %[[ARG0]] : tensor<4x6x16x32xi8> -func @trivial_subtensor_insert(%arg0 : tensor<4x6x16x32xi8>, %arg1 : tensor<4x6x16x32xi8>) -> tensor<4x6x16x32xi8> { - %0 = subtensor_insert %arg0 into %arg1[0, 0, 0, 0] [4, 6, 16, 32] [1, 1, 1, 1] : tensor<4x6x16x32xi8> into tensor<4x6x16x32xi8> - return %0 : tensor<4x6x16x32xi8> -} - -// ----- - -// CHECK-LABEL: func @rank_reducing_tensor_of_cast -// CHECK-SAME: %[[ARG0:.[a-z0-9A-Z_]+]]: tensor<4x6x16x32xi8> -// CHECK: %[[S:.+]] = subtensor %arg0[0, 1, 0] [1, 1, 16] [1, 1, 1] : tensor<4x6x16x32xi8> to tensor<16x32xi8> -// Tensor cast is moved after subtensor and then gets canonicalized away. -// CHECK-NOT: tensor.cast -// CHECK: return %[[S]] : tensor<16x32xi8> -func @rank_reducing_tensor_of_cast(%arg : tensor<4x6x16x32xi8>) -> tensor<16x32xi8> { - %0 = tensor.cast %arg : tensor<4x6x16x32xi8> to tensor - %1 = subtensor %0[0, 1, 0] [1, 1, 16] [1, 1, 1] : tensor to tensor<16x32xi8> - return %1 : tensor<16x32xi8> -} - -// ----- - -// CHECK-LABEL: func @rank_reducing_subtensor_insert_of_cast -// CHECK-SAME: %[[A:.[a-z0-9A-Z_]+]]: tensor<16x32xi8> -// CHECK-SAME: %[[B:.[a-z0-9A-Z_]+]]: tensor<4x6x16x32xi8> -// CHECK: %[[S:.+]] = subtensor_insert %[[A]] into %[[B]][0, 1, 0] [1, 1, 16] [1, 1, 1] : tensor<16x32xi8> into tensor<4x6x16x32xi8> -// Tensor cast is folded away. -// CHECK-NOT: tensor.cast -// CHECK: return %[[S]] : tensor<4x6x16x32xi8> -func @rank_reducing_subtensor_insert_of_cast(%a : tensor<16x32xi8>, %b : tensor<4x6x16x32xi8>) -> tensor<4x6x16x32xi8> { - %cast = tensor.cast %a : tensor<16x32xi8> to tensor - %res = subtensor_insert %cast into %b[0, 1, 0] [1, 1, 16] [1, 1, 1] : tensor into tensor<4x6x16x32xi8> - return %res : tensor<4x6x16x32xi8> -} - -// ----- - -func @subtensor_insert_canonicalize(%arg0 : tensor, %arg1 : index, - %arg2 : index, %arg3 : tensor) -> tensor -{ - %c0 = constant 0 : index - %c1 = constant 1 : index - %c4 = constant 4 : index - %0 = subtensor_insert %arg0 into %arg3[%c0, %arg1, %c1] [%c4, %c1, %arg2] [%c1, %c1, %c1] : tensor into tensor - return %0 : tensor -} -// CHECK-LABEL: func @subtensor_insert_canonicalize -// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor -// CHECK: %[[RESULT:.+]] = subtensor_insert %[[ARG0]] -// CHECK-SAME: [0, %{{.+}}, 1] [4, 1, %{{.+}}] [1, 1, 1] -// CHECK-SAME: : tensor into tensor -// CHEKC: return %[[RESULT]] - -// ----- - -func @subtensor_to_subtensor_insert_canonicalize(%arg0 : tensor, %arg1 : index, - %arg2 : index, %arg3 : tensor) -> tensor -{ - %c0 = constant 0 : index - %c1 = constant 1 : index - %c4 = constant 4 : index - %0 = subtensor %arg0[%c0, %arg1, %c1] [%c4, %c1, %arg2] [%c1, %c1, %c1] : tensor to tensor - %1 = subtensor_insert %0 into %arg3[%c0, %arg1, %c1] [%c4, %c1, %arg2] [%c1, %c1, %c1] : tensor into tensor - return %1 : tensor -} -// CHECK-LABEL: func @subtensor_to_subtensor_insert_canonicalize -// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor -// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: tensor -// CHECK: %[[SUBTENSOR:.+]] = subtensor %[[ARG0]] -// CHECK-SAME: [0, %{{.+}}, 1] [4, 1, %{{.+}} [1, 1, 1] -// CHECK-SAME: : tensor to tensor<4x1x?xf32> -// CHECK: %[[RESULT:.+]] = subtensor_insert %[[SUBTENSOR]] -// CHECK-SAME: [0, %{{.+}}, 1] [4, 1, %{{.+}}] [1, 1, 1] -// CHECK-SAME: : tensor<4x1x?xf32> into tensor -// CHEKC: return %[[RESULT]] - -// ----- - -func @rank_reducing_subtensor_insert_canonicalize(%arg0 : tensor, %arg1 : index, - %arg2 : index, %arg3 : tensor) -> tensor -{ - %c0 = constant 0 : index - %c1 = constant 1 : index - %c4 = constant 4 : index - %0 = subtensor_insert %arg0 into %arg3[%c0, %arg1, %c1] [%c4, 1, %arg2] [%c1, %c1, %c1] : tensor into tensor - return %0 : tensor -} -// CHECK-LABEL: func @rank_reducing_subtensor_insert_canonicalize -// CHECK-SAME: %[[ARG0:.+]]: tensor -// CHECK: %[[RESULT:.+]] = subtensor_insert %[[ARG0]] -// CHECK-SAME: [0, %{{.+}}, 1] [4, 1, %{{.+}}] [1, 1, 1] -// CHECK-SAME: : tensor into tensor -// CHEKC: return %[[RESULT]] - -// ----- - -func @rank_reducing_subtensor_to_subtensor_insert_canonicalize(%arg0 : tensor, %arg1 : index, - %arg2 : index, %arg3 : tensor) -> tensor -{ - %c0 = constant 0 : index - %c1 = constant 1 : index - %c4 = constant 4 : index - %0 = subtensor %arg0[%c0, %arg1, %c1] [%c4, 1, %arg2] [%c1, %c1, %c1] : tensor to tensor - %1 = subtensor_insert %0 into %arg3[%c0, %arg1, %c1] [%c4, 1, %arg2] [%c1, %c1, %c1] : tensor into tensor - return %1 : tensor -} -// CHECK-LABEL: func @rank_reducing_subtensor_to_subtensor_insert_canonicalize -// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor -// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: tensor -// CHECK: %[[SUBTENSOR:.+]] = subtensor %[[ARG0]] -// CHECK-SAME: [0, %{{.+}}, 1] [4, 1, %{{.+}}] [1, 1, 1] -// CHECK-SAME: : tensor to tensor<4x?xf32> -// CHECK: %[[RESULT:.+]] = subtensor_insert %[[SUBTENSOR]] into %[[ARG3]] -// CHECK-SAME: [0, %{{.+}}, 1] [4, 1, %{{.+}}] [1, 1, 1] -// CHECK-SAME: : tensor<4x?xf32> into tensor -// CHEKC: return %[[RESULT]] - -// ----- - -func @subtensor_insert_propagate_dest_cast(%arg0 : tensor<2x?xi32>, %arg1 : tensor, - %arg2 : index, %arg3 : index) -> tensor { - %c0 = constant 0 : index - %c1 = constant 1 : index - %c2 = constant 2 : index - %c8 = constant 8 : index - %0 = memref.dim %arg0, %c1 : tensor<2x?xi32> - %1 = tensor.extract %arg1[] : tensor - %2 = tensor.generate %arg2, %c8 { - ^bb0(%arg4: index, %arg5: index): - tensor.yield %1 : i32 - } : tensor - %3 = subtensor_insert %arg0 into %2[%c0, %arg3] [%c2, %0] [%c1, %c1] : tensor<2x?xi32> into tensor - return %3 : tensor -} -// CHECK-LABEL: func @subtensor_insert_propagate_dest_cast -// CHECK: %[[UPDATED:.+]] = subtensor_insert %{{.+}} into %{{.+}}[0, %{{.+}}] [2, %{{.+}}] [1, 1] -// CHECK-SAME: tensor<2x?xi32> into tensor -// CHECK: %[[CAST:.+]] = tensor.cast %[[UPDATED]] -// CHECK: return %[[CAST]] - -// ----- - -func @subtensor_insert_output_dest_canonicalize(%arg0 : tensor<2x3xi32>, %arg1 : tensor) -> tensor<3x9xi32> { - %c0 = constant 0 : index - %c1 = constant 1 : index - %c2 = constant 2 : index - %c9 = constant 9 : index - %c3 = constant 3 : index - %2 = tensor.extract %arg1[] : tensor - %4 = tensor.generate %c3, %c9 { - ^bb0(%arg2: index, %arg3: index): - tensor.yield %2 : i32 - } : tensor - %5 = subtensor_insert %arg0 into %4[%c0, %c1] [%c2, %c3] [1, 1] : tensor<2x3xi32> into tensor - %6 = tensor.cast %5 : tensor to tensor<3x9xi32> - return %6 : tensor<3x9xi32> -} -// CHECK-LABEL: func @subtensor_insert_output_dest_canonicalize -// CHECK-SAME: %[[ARG0:[a-zA-z0-9_]+]]: tensor<2x3xi32> -// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor -// CHECK: %[[PAD:.+]] = tensor.extract %[[ARG1]] -// CHECK: %[[GENERATE:.+]] = tensor.generate -// CHECK: %[[RESULT:.+]] = subtensor_insert %[[ARG0]] into %[[GENERATE]] -// CHECK: return %[[RESULT]] - -// ----- - // CHECK-LABEL: @select_same_val // CHECK: return %arg1 func @select_same_val(%arg0: i1, %arg1: i64) -> i64 { diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir --- a/mlir/test/Dialect/Tensor/canonicalize.mlir +++ b/mlir/test/Dialect/Tensor/canonicalize.mlir @@ -263,3 +263,222 @@ %tensor = tensor.from_elements %c1, %c2, %c1 : tensor<3xindex> return %tensor : tensor<3xindex> } + +// ----- + +func @subtensor_canonicalize(%arg0 : tensor, %arg1 : index, + %arg2 : index) -> tensor +{ + %c0 = constant 0 : index + %c1 = constant 1 : index + %c4 = constant 4 : index + %0 = tensor.subtensor %arg0[%c0, %arg1, %c1] [%c4, %c1, %arg2] [%c1, %c1, %c1] : tensor to tensor + return %0 : tensor +} +// CHECK-LABEL: func @subtensor_canonicalize +// CHECK-SAME: %[[ARG0:.+]]: tensor +// CHECK: %[[SUBTENSOR:.+]] = tensor.subtensor %[[ARG0]][0, %{{[a-zA-Z0-9_]+}}, 1] +// CHECK-SAME: [4, 1, %{{[a-zA-Z0-9_]+}}] [1, 1, 1] +// CHECK-SAME: : tensor to tensor<4x1x?xf32> +// CHECK: %[[RESULT:.+]] = tensor.cast %[[SUBTENSOR]] +// CHEKC: return %[[RESULT]] + +// ----- + +func @rank_reducing_subtensor_canonicalize(%arg0 : tensor, %arg1 : index, + %arg2 : index) -> tensor +{ + %c0 = constant 0 : index + %c1 = constant 1 : index + %c4 = constant 4 : index + %0 = tensor.subtensor %arg0[%c0, %arg1, %c1] [%c4, 1, %arg2] [%c1, %c1, %c1] : tensor to tensor + return %0 : tensor +} +// CHECK-LABEL: func @rank_reducing_subtensor_canonicalize +// CHECK-SAME: %[[ARG0:.+]]: tensor +// CHECK: %[[SUBTENSOR:.+]] = tensor.subtensor %[[ARG0]][0, %{{[a-zA-Z0-9_]+}}, 1] +// CHECK-SAME: [4, 1, %{{[a-zA-Z0-9_]+}}] [1, 1, 1] +// CHECK-SAME: : tensor to tensor<4x?xf32> +// CHECK: %[[RESULT:.+]] = tensor.cast %[[SUBTENSOR]] +// CHEKC: return %[[RESULT]] + +// ----- + +// CHECK-LABEL: func @trivial_subtensor +// CHECK-SAME: %[[ARG0:.[a-z0-9A-Z_]+]]: tensor<4x6x16x32xi8> +// CHECK-NOT: tensor.subtensor +// CHECK: return %[[ARG0]] : tensor<4x6x16x32xi8> +func @trivial_subtensor(%arg0 : tensor<4x6x16x32xi8>) -> tensor<4x6x16x32xi8> { + %0 = tensor.subtensor %arg0[0, 0, 0, 0] [4, 6, 16, 32] [1, 1, 1, 1] : tensor<4x6x16x32xi8> to tensor<4x6x16x32xi8> + return %0 : tensor<4x6x16x32xi8> +} + +// ----- + +// CHECK-LABEL: func @trivial_subtensor_insert +// CHECK-SAME: %[[ARG0:.[a-z0-9A-Z_]+]]: tensor<4x6x16x32xi8> +// CHECK-NOT: tensor.subtensor +// CHECK: return %[[ARG0]] : tensor<4x6x16x32xi8> +func @trivial_subtensor_insert(%arg0 : tensor<4x6x16x32xi8>, %arg1 : tensor<4x6x16x32xi8>) -> tensor<4x6x16x32xi8> { + %0 = tensor.subtensor_insert %arg0 into %arg1[0, 0, 0, 0] [4, 6, 16, 32] [1, 1, 1, 1] : tensor<4x6x16x32xi8> into tensor<4x6x16x32xi8> + return %0 : tensor<4x6x16x32xi8> +} + +// ----- + +// CHECK-LABEL: func @rank_reducing_tensor_of_cast +// CHECK-SAME: %[[ARG0:.[a-z0-9A-Z_]+]]: tensor<4x6x16x32xi8> +// CHECK: %[[S:.+]] = tensor.subtensor %arg0[0, 1, 0] [1, 1, 16] [1, 1, 1] : tensor<4x6x16x32xi8> to tensor<16x32xi8> +// Tensor cast is moved after subtensor and then gets canonicalized away. +// CHECK-NOT: tensor.cast +// CHECK: return %[[S]] : tensor<16x32xi8> +func @rank_reducing_tensor_of_cast(%arg : tensor<4x6x16x32xi8>) -> tensor<16x32xi8> { + %0 = tensor.cast %arg : tensor<4x6x16x32xi8> to tensor + %1 = tensor.subtensor %0[0, 1, 0] [1, 1, 16] [1, 1, 1] : tensor to tensor<16x32xi8> + return %1 : tensor<16x32xi8> +} + +// ----- + +// CHECK-LABEL: func @rank_reducing_subtensor_insert_of_cast +// CHECK-SAME: %[[A:.[a-z0-9A-Z_]+]]: tensor<16x32xi8> +// CHECK-SAME: %[[B:.[a-z0-9A-Z_]+]]: tensor<4x6x16x32xi8> +// CHECK: %[[S:.+]] = tensor.subtensor_insert %[[A]] into %[[B]][0, 1, 0] [1, 1, 16] [1, 1, 1] : tensor<16x32xi8> into tensor<4x6x16x32xi8> +// Tensor cast is folded away. +// CHECK-NOT: tensor.cast +// CHECK: return %[[S]] : tensor<4x6x16x32xi8> +func @rank_reducing_subtensor_insert_of_cast(%a : tensor<16x32xi8>, %b : tensor<4x6x16x32xi8>) -> tensor<4x6x16x32xi8> { + %cast = tensor.cast %a : tensor<16x32xi8> to tensor + %res = tensor.subtensor_insert %cast into %b[0, 1, 0] [1, 1, 16] [1, 1, 1] : tensor into tensor<4x6x16x32xi8> + return %res : tensor<4x6x16x32xi8> +} + +// ----- + +func @subtensor_insert_canonicalize(%arg0 : tensor, %arg1 : index, + %arg2 : index, %arg3 : tensor) -> tensor +{ + %c0 = constant 0 : index + %c1 = constant 1 : index + %c4 = constant 4 : index + %0 = tensor.subtensor_insert %arg0 into %arg3[%c0, %arg1, %c1] [%c4, %c1, %arg2] [%c1, %c1, %c1] : tensor into tensor + return %0 : tensor +} +// CHECK-LABEL: func @subtensor_insert_canonicalize +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor +// CHECK: %[[RESULT:.+]] = tensor.subtensor_insert %[[ARG0]] +// CHECK-SAME: [0, %{{.+}}, 1] [4, 1, %{{.+}}] [1, 1, 1] +// CHECK-SAME: : tensor into tensor +// CHEKC: return %[[RESULT]] + +// ----- + +func @subtensor_to_subtensor_insert_canonicalize(%arg0 : tensor, %arg1 : index, + %arg2 : index, %arg3 : tensor) -> tensor +{ + %c0 = constant 0 : index + %c1 = constant 1 : index + %c4 = constant 4 : index + %0 = tensor.subtensor %arg0[%c0, %arg1, %c1] [%c4, %c1, %arg2] [%c1, %c1, %c1] : tensor to tensor + %1 = tensor.subtensor_insert %0 into %arg3[%c0, %arg1, %c1] [%c4, %c1, %arg2] [%c1, %c1, %c1] : tensor into tensor + return %1 : tensor +} +// CHECK-LABEL: func @subtensor_to_subtensor_insert_canonicalize +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor +// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: tensor +// CHECK: %[[SUBTENSOR:.+]] = tensor.subtensor %[[ARG0]] +// CHECK-SAME: [0, %{{.+}}, 1] [4, 1, %{{.+}} [1, 1, 1] +// CHECK-SAME: : tensor to tensor<4x1x?xf32> +// CHECK: %[[RESULT:.+]] = tensor.subtensor_insert %[[SUBTENSOR]] +// CHECK-SAME: [0, %{{.+}}, 1] [4, 1, %{{.+}}] [1, 1, 1] +// CHECK-SAME: : tensor<4x1x?xf32> into tensor +// CHEKC: return %[[RESULT]] + +// ----- + +func @rank_reducing_subtensor_insert_canonicalize(%arg0 : tensor, %arg1 : index, + %arg2 : index, %arg3 : tensor) -> tensor +{ + %c0 = constant 0 : index + %c1 = constant 1 : index + %c4 = constant 4 : index + %0 = tensor.subtensor_insert %arg0 into %arg3[%c0, %arg1, %c1] [%c4, 1, %arg2] [%c1, %c1, %c1] : tensor into tensor + return %0 : tensor +} +// CHECK-LABEL: func @rank_reducing_subtensor_insert_canonicalize +// CHECK-SAME: %[[ARG0:.+]]: tensor +// CHECK: %[[RESULT:.+]] = tensor.subtensor_insert %[[ARG0]] +// CHECK-SAME: [0, %{{.+}}, 1] [4, 1, %{{.+}}] [1, 1, 1] +// CHECK-SAME: : tensor into tensor +// CHEKC: return %[[RESULT]] + +// ----- + +func @rank_reducing_subtensor_to_subtensor_insert_canonicalize(%arg0 : tensor, %arg1 : index, + %arg2 : index, %arg3 : tensor) -> tensor +{ + %c0 = constant 0 : index + %c1 = constant 1 : index + %c4 = constant 4 : index + %0 = tensor.subtensor %arg0[%c0, %arg1, %c1] [%c4, 1, %arg2] [%c1, %c1, %c1] : tensor to tensor + %1 = tensor.subtensor_insert %0 into %arg3[%c0, %arg1, %c1] [%c4, 1, %arg2] [%c1, %c1, %c1] : tensor into tensor + return %1 : tensor +} +// CHECK-LABEL: func @rank_reducing_subtensor_to_subtensor_insert_canonicalize +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor +// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: tensor +// CHECK: %[[SUBTENSOR:.+]] = tensor.subtensor %[[ARG0]] +// CHECK-SAME: [0, %{{.+}}, 1] [4, 1, %{{.+}}] [1, 1, 1] +// CHECK-SAME: : tensor to tensor<4x?xf32> +// CHECK: %[[RESULT:.+]] = tensor.subtensor_insert %[[SUBTENSOR]] into %[[ARG3]] +// CHECK-SAME: [0, %{{.+}}, 1] [4, 1, %{{.+}}] [1, 1, 1] +// CHECK-SAME: : tensor<4x?xf32> into tensor +// CHEKC: return %[[RESULT]] + +// ----- + +func @subtensor_insert_propagate_dest_cast(%arg0 : tensor<2x?xi32>, %arg1 : tensor, + %arg2 : index, %arg3 : index) -> tensor { + %c0 = constant 0 : index + %c1 = constant 1 : index + %c2 = constant 2 : index + %c8 = constant 8 : index + %0 = memref.dim %arg0, %c1 : tensor<2x?xi32> + %1 = tensor.extract %arg1[] : tensor + %2 = tensor.generate %arg2, %c8 { + ^bb0(%arg4: index, %arg5: index): + tensor.yield %1 : i32 + } : tensor + %3 = tensor.subtensor_insert %arg0 into %2[%c0, %arg3] [%c2, %0] [%c1, %c1] : tensor<2x?xi32> into tensor + return %3 : tensor +} +// CHECK-LABEL: func @subtensor_insert_propagate_dest_cast +// CHECK: %[[UPDATED:.+]] = tensor.subtensor_insert %{{.+}} into %{{.+}}[0, %{{.+}}] [2, %{{.+}}] [1, 1] +// CHECK-SAME: tensor<2x?xi32> into tensor +// CHECK: %[[CAST:.+]] = tensor.cast %[[UPDATED]] +// CHECK: return %[[CAST]] + +// ----- + +func @subtensor_insert_output_dest_canonicalize(%arg0 : tensor<2x3xi32>, %arg1 : tensor) -> tensor<3x9xi32> { + %c0 = constant 0 : index + %c1 = constant 1 : index + %c2 = constant 2 : index + %c9 = constant 9 : index + %c3 = constant 3 : index + %2 = tensor.extract %arg1[] : tensor + %4 = tensor.generate %c3, %c9 { + ^bb0(%arg2: index, %arg3: index): + tensor.yield %2 : i32 + } : tensor + %5 = tensor.subtensor_insert %arg0 into %4[%c0, %c1] [%c2, %c3] [1, 1] : tensor<2x3xi32> into tensor + %6 = tensor.cast %5 : tensor to tensor<3x9xi32> + return %6 : tensor<3x9xi32> +} +// CHECK-LABEL: func @subtensor_insert_output_dest_canonicalize +// CHECK-SAME: %[[ARG0:[a-zA-z0-9_]+]]: tensor<2x3xi32> +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor +// CHECK: %[[PAD:.+]] = tensor.extract %[[ARG1]] +// CHECK: %[[GENERATE:.+]] = tensor.generate +// CHECK: %[[RESULT:.+]] = tensor.subtensor_insert %[[ARG0]] into %[[GENERATE]] +// CHECK: return %[[RESULT]] 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 @@ -830,19 +830,19 @@ %c0 = constant 0 : index %c1 = constant 1 : index - // CHECK: subtensor + // CHECK: tensor.subtensor // CHECK-SAME: tensor<8x16x4xf32> to tensor - %1 = subtensor %t[%c0, %c0, %c0][%idx, %idx, %idx][%c1, %c1, %c1] + %1 = tensor.subtensor %t[%c0, %c0, %c0][%idx, %idx, %idx][%c1, %c1, %c1] : tensor<8x16x4xf32> to tensor - // CHECK: subtensor + // CHECK: tensor.subtensor // CHECK-SAME: tensor<8x16x4xf32> to tensor<4x4x4xf32> - %2 = subtensor %t[0, 2, 0][4, 4, 4][1, 1, 1] + %2 = tensor.subtensor %t[0, 2, 0][4, 4, 4][1, 1, 1] : tensor<8x16x4xf32> to tensor<4x4x4xf32> - // CHECK: subtensor + // CHECK: tensor.subtensor // CHECK-SAME: tensor<8x16x4xf32> to tensor<4x4xf32> - %3 = subtensor %t[0, 2, 0][4, 1, 4][1, 1, 1] + %3 = tensor.subtensor %t[0, 2, 0][4, 1, 4][1, 1, 1] : tensor<8x16x4xf32> to tensor<4x4xf32> return @@ -857,19 +857,19 @@ %c0 = constant 0 : index %c1 = constant 1 : index - // CHECK: subtensor_insert + // CHECK: tensor.subtensor_insert // CHECK-SAME: tensor<8x16x4xf32> into tensor<16x32x8xf32> - %1 = subtensor_insert %t into %t2[%c0, %c0, %c0][%idx, %idx, %idx][%c1, %c1, %c1] + %1 = tensor.subtensor_insert %t into %t2[%c0, %c0, %c0][%idx, %idx, %idx][%c1, %c1, %c1] : tensor<8x16x4xf32> into tensor<16x32x8xf32> - // CHECK: subtensor_insert + // CHECK: tensor.subtensor_insert // CHECK-SAME: tensor<8x16x4xf32> into tensor<16x32x8xf32> - %2 = subtensor_insert %t into %t2[%c0, %idx, %c0][%idx, 4, %idx][%c1, 1, %c1] + %2 = tensor.subtensor_insert %t into %t2[%c0, %idx, %c0][%idx, 4, %idx][%c1, 1, %c1] : tensor<8x16x4xf32> into tensor<16x32x8xf32> - // CHECK: subtensor_insert + // CHECK: tensor.subtensor_insert // CHECK-SAME: tensor<4x4xf32> into tensor<8x16x4xf32> - %3 = subtensor_insert %t3 into %t[0, 2, 0][4, 1, 4][1, 1, 1] + %3 = tensor.subtensor_insert %t3 into %t[0, 2, 0][4, 1, 4][1, 1, 1] : tensor<4x4xf32> into tensor<8x16x4xf32> return 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 @@ -1216,7 +1216,7 @@ func @subtensor_wrong_dynamic_type(%t: tensor<8x16x4xf32>, %idx : index) { // expected-error @+1 {{expected result type to be 'tensor<4x4x4xf32>' or a rank-reduced version. (mismatch of result sizes)}} - %0 = subtensor %t[0, 2, 0][4, 4, 4][1, 1, 1] + %0 = tensor.subtensor %t[0, 2, 0][4, 4, 4][1, 1, 1] : tensor<8x16x4xf32> to tensor return @@ -1226,7 +1226,7 @@ func @subtensor_wrong_static_type(%t: tensor<8x16x4xf32>, %idx : index) { // expected-error @+1 {{expected result type to be 'tensor' or a rank-reduced version. (mismatch of result sizes)}} - %0 = subtensor %t[0, 0, 0][%idx, 3, %idx][1, 1, 1] + %0 = tensor.subtensor %t[0, 0, 0][%idx, 3, %idx][1, 1, 1] : tensor<8x16x4xf32> to tensor<4x4x4xf32> return diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/test-subtensor-insert-multiple-uses.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/test-subtensor-insert-multiple-uses.mlir --- a/mlir/test/Integration/Dialect/Linalg/CPU/test-subtensor-insert-multiple-uses.mlir +++ b/mlir/test/Integration/Dialect/Linalg/CPU/test-subtensor-insert-multiple-uses.mlir @@ -14,8 +14,8 @@ // value `%const`. This can easily cause bugs if at the memref level // we attempt to write in-place into the memref that %const has been // converted into. - %inserted_at_position_0 = subtensor_insert %insert_val into %const[0][1][1] : tensor<1xf32> into tensor<2xf32> - %inserted_at_position_1 = subtensor_insert %insert_val into %const[1][1][1] : tensor<1xf32> into tensor<2xf32> + %inserted_at_position_0 = tensor.subtensor_insert %insert_val into %const[0][1][1] : tensor<1xf32> into tensor<2xf32> + %inserted_at_position_1 = tensor.subtensor_insert %insert_val into %const[1][1][1] : tensor<1xf32> into tensor<2xf32> %unranked_at_position_0 = tensor.cast %inserted_at_position_0 : tensor<2xf32> to tensor<*xf32> call @print_memref_f32(%unranked_at_position_0) : (tensor<*xf32>) -> () diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/test-subtensor-insert.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/test-subtensor-insert.mlir --- a/mlir/test/Integration/Dialect/Linalg/CPU/test-subtensor-insert.mlir +++ b/mlir/test/Integration/Dialect/Linalg/CPU/test-subtensor-insert.mlir @@ -9,7 +9,7 @@ func @main() { %const = constant dense<10.0> : tensor<2xf32> %insert_val = constant dense<20.0> : tensor<1xf32> - %inserted = subtensor_insert %insert_val into %const[0][1][1] : tensor<1xf32> into tensor<2xf32> + %inserted = tensor.subtensor_insert %insert_val into %const[0][1][1] : tensor<1xf32> into tensor<2xf32> %unranked = tensor.cast %inserted : tensor<2xf32> to tensor<*xf32> call @print_memref_f32(%unranked) : (tensor<*xf32>) -> () 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 @@ -1076,18 +1076,18 @@ %c7 = constant 7 : index %c11 = constant 11 : index - // CHECK: subtensor %{{.*}}[0, 0, 0] [7, 11, 2] [1, 1, 1] : + // CHECK: tensor.subtensor %{{.*}}[0, 0, 0] [7, 11, 2] [1, 1, 1] : // CHECK-SAME: tensor<8x16x4xf32> to tensor<7x11x2xf32> // tensor.cast gets folded away in consumer. // CHECK-NOT: tensor.cast - %1 = subtensor %t[%c0, %c0, %c0] [%c7, %c11, %c2] [%c1, %c1, %c1] + %1 = tensor.subtensor %t[%c0, %c0, %c0] [%c7, %c11, %c2] [%c1, %c1, %c1] : tensor<8x16x4xf32> to tensor // Test: subtensor with one dynamic operand can also be folded. - // CHECK: subtensor %{{.*}}[0, 0, 0] [2, %[[ARG0]], 2] [1, 1, 1] : + // CHECK: tensor.subtensor %{{.*}}[0, 0, 0] [2, %[[ARG0]], 2] [1, 1, 1] : // CHECK-SAME: tensor<7x11x2xf32> to tensor<2x?x2xf32> // CHECK: tensor.cast %{{.*}} : tensor<2x?x2xf32> to tensor - %2 = subtensor %1[%c0, %c0, %c0] [%c2, %arg0, %c2] [%c1, %c1, %c1] + %2 = tensor.subtensor %1[%c0, %c0, %c0] [%c2, %arg0, %c2] [%c1, %c1, %c1] : tensor to tensor return %2 : tensor