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 @@ -191,14 +191,17 @@ let builders = [ // Build with fully static sizes. - OpBuilder<(ins "ArrayRef":$staticShape, "Type":$elementType)>, + OpBuilder<(ins "ArrayRef":$staticShape, "Type":$elementType, + CArg<"Attribute", "{}">:$encoding)>, // Build with mixed static/dynamic sizes. OpBuilder<(ins "ArrayRef":$staticShape, "Type":$elementType, - "ValueRange":$dynamicSizes)>, + "ValueRange":$dynamicSizes, + CArg<"Attribute", "{}">:$encoding)>, // Build with mixed static/dynamic sizes. - OpBuilder<(ins "ArrayRef":$sizes, "Type":$elementType)> + OpBuilder<(ins "ArrayRef":$sizes, "Type":$elementType, + CArg<"Attribute", "{}">:$encoding)> ]; let hasCanonicalizer = 1; 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 @@ -497,27 +497,29 @@ //===----------------------------------------------------------------------===// void EmptyOp::build(OpBuilder &builder, OperationState &result, - ArrayRef staticShape, Type elementType) { + ArrayRef staticShape, Type elementType, + Attribute encoding) { assert(all_of(staticShape, [](int64_t sz) { return !ShapedType::isDynamic(sz); }) && "expected only static sizes"); - build(builder, result, staticShape, elementType, {}); + build(builder, result, staticShape, elementType, ValueRange{}, encoding); } void EmptyOp::build(OpBuilder &builder, OperationState &result, ArrayRef staticShape, Type elementType, - ValueRange dynamicSizes) { - auto tensorType = RankedTensorType::get(staticShape, elementType); + ValueRange dynamicSizes, Attribute encoding) { + auto tensorType = RankedTensorType::get(staticShape, elementType, encoding); build(builder, result, tensorType, dynamicSizes); } void EmptyOp::build(OpBuilder &builder, OperationState &result, - ArrayRef sizes, Type elementType) { + ArrayRef sizes, Type elementType, + Attribute encoding) { SmallVector staticShape; SmallVector dynamicSizes; dispatchIndexOpFoldResults(sizes, dynamicSizes, staticShape, ShapedType::kDynamicSize); - build(builder, result, staticShape, elementType, dynamicSizes); + build(builder, result, staticShape, elementType, dynamicSizes, encoding); } LogicalResult EmptyOp::verify() { diff --git a/mlir/test/Dialect/Tensor/ops.mlir b/mlir/test/Dialect/Tensor/ops.mlir --- a/mlir/test/Dialect/Tensor/ops.mlir +++ b/mlir/test/Dialect/Tensor/ops.mlir @@ -21,6 +21,15 @@ return %0 : tensor<5x?x6xf32> } +// CHECK-LABEL: func @empty_with_encoding( +// CHECK-SAME: %[[sz:.*]]: index +func.func @empty_with_encoding(%sz: index) -> tensor<5x?x6xf32, "foo"> { + // CHECK: tensor.empty(%[[sz]]) : tensor<5x?x6xf32, "foo"> + %0 = tensor.empty(%sz) : tensor<5x?x6xf32, "foo"> + return %0 : tensor<5x?x6xf32, "foo"> +} + + // CHECK-LABEL: func @extract( // CHECK-SAME: %[[TENSOR:.*]]: tensor, // CHECK-SAME: %[[INDEX:.*]]: index) {