Support for TF.Gather including intermediate dialect work (#4292)
Gather operations added to the strings and tf_strings dialect with
lowerings from TensorFlow to tf_strings and tf_string to strings.
This includes tests for each and an end-to-end test. Right now the
python test does not validate values (missing conversion from
StringTensor to Numpy equivalent) but prints instead. Close enough
for current support.
diff --git a/integrations/tensorflow/compiler/Passes.cpp b/integrations/tensorflow/compiler/Passes.cpp
index e208e99..3b7593a 100644
--- a/integrations/tensorflow/compiler/Passes.cpp
+++ b/integrations/tensorflow/compiler/Passes.cpp
@@ -72,6 +72,12 @@
pm.addPass(createCanonicalizerPass());
//----------------------------------------------------------------------------
+ // Lowering module specific TF behavior to intermediate dialects.
+ //----------------------------------------------------------------------------
+ pm.addPass(tf_strings::createConvertTFToTFStringsPass());
+ pm.addPass(tf_tensorlist::createConvertTFToTFTensorListPass());
+
+ //----------------------------------------------------------------------------
// Legalize to XLA
//----------------------------------------------------------------------------
pm.addPass(createConvertToMHLOPass());
@@ -93,11 +99,8 @@
pm.addPass(createCanonicalizerPass());
//----------------------------------------------------------------------------
- // Lowering module specific TF behavior to custom dialects.
+ // Lowering intermediate dialects to module specific dialects.
//----------------------------------------------------------------------------
- pm.addPass(tf_strings::createConvertTFToTFStringsPass());
- pm.addPass(tf_tensorlist::createConvertTFToTFTensorListPass());
- pm.addPass(createCanonicalizerPass());
pm.addPass(tf_strings::createConvertTFStringsToStringsPass());
pm.addPass(tf_tensorlist::createConvertTFTensorListToTensorListPass());
diff --git a/integrations/tensorflow/compiler/dialect/tf_strings/conversion/convert_tf_strings_to_strings.cc b/integrations/tensorflow/compiler/dialect/tf_strings/conversion/convert_tf_strings_to_strings.cc
index 7217eb7..e043967 100644
--- a/integrations/tensorflow/compiler/dialect/tf_strings/conversion/convert_tf_strings_to_strings.cc
+++ b/integrations/tensorflow/compiler/dialect/tf_strings/conversion/convert_tf_strings_to_strings.cc
@@ -80,6 +80,9 @@
patterns.insert<OpConversion<tf_strings::ToStringTensorOp,
iree_compiler::IREE::Strings::ToStringTensorOp>>(
context);
+ patterns.insert<OpConversion<tf_strings::GatherOp,
+ iree_compiler::IREE::Strings::GatherOp>>(
+ context);
patterns.insert<
OpConversion<tf_strings::StringTensorToStringOp,
iree_compiler::IREE::Strings::StringTensorToStringOp>>(
diff --git a/integrations/tensorflow/compiler/dialect/tf_strings/conversion/convert_tf_to_tf_strings.cc b/integrations/tensorflow/compiler/dialect/tf_strings/conversion/convert_tf_to_tf_strings.cc
index e3baa11..4fe4d7e 100644
--- a/integrations/tensorflow/compiler/dialect/tf_strings/conversion/convert_tf_to_tf_strings.cc
+++ b/integrations/tensorflow/compiler/dialect/tf_strings/conversion/convert_tf_to_tf_strings.cc
@@ -23,6 +23,7 @@
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Dialect.h"
+#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassRegistry.h"
@@ -85,6 +86,44 @@
}
};
+struct GatherV2OpLowering : public OpRewritePattern<TF::GatherV2Op> {
+ using OpRewritePattern<TF::GatherV2Op>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(TF::GatherV2Op op,
+ PatternRewriter &rewriter) const override {
+ auto tensor = op.params();
+ auto tensorTy = tensor.getType().dyn_cast<RankedTensorType>();
+ if (!tensorTy || !tensorTy.getElementType().isa<TF::StringType>()) {
+ return failure();
+ }
+
+ DenseIntElementsAttr axis;
+ if (!matchPattern(op.axis(), m_Constant(&axis))) {
+ return failure();
+ }
+
+ if (axis.getType().cast<ShapedType>().getRank() != 0) {
+ return failure();
+ }
+
+ auto axisValue = axis.getValue<IntegerAttr>({});
+ auto axisInt = axisValue.getValue().getZExtValue();
+
+ if (axisInt != tensorTy.getRank() - 1) {
+ return failure();
+ }
+
+ auto resultTy = op.getType().cast<ShapedType>();
+ rewriter.replaceOpWithNewOp<tf_strings::GatherOp>(
+ op,
+ RankedTensorType::get(resultTy.getShape(),
+ tf_strings::StringType::get(op.getContext())),
+ tensor, op.indices());
+
+ return success();
+ }
+};
+
class ConvertTFToTFStringsPass
: public ConversionPass<ConvertTFToTFStringsPass, StringTypeConverter> {
public:
@@ -108,6 +147,7 @@
void populateTFToTFStringsPatterns(MLIRContext *ctx,
OwningRewritePatternList &patterns) {
populateWithGenerated(ctx, patterns);
+ patterns.insert<GatherV2OpLowering>(ctx);
patterns.insert<StringFormatOpLowering>(ctx);
}
diff --git a/integrations/tensorflow/compiler/dialect/tf_strings/conversion/convert_tf_to_tf_strings.td b/integrations/tensorflow/compiler/dialect/tf_strings/conversion/convert_tf_to_tf_strings.td
index 83c9d6f..f20380d 100644
--- a/integrations/tensorflow/compiler/dialect/tf_strings/conversion/convert_tf_to_tf_strings.td
+++ b/integrations/tensorflow/compiler/dialect/tf_strings/conversion/convert_tf_to_tf_strings.td
@@ -17,6 +17,9 @@
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
include "integrations/tensorflow/compiler/dialect/tf_strings/ir/ops.td"
+def : Pat<(TF_IdentityOp TF_StrTensor:$value),
+ (replaceWithValue $value)>;
+
def : Pat<(TF_PrintV2Op TF_Str:$value, $unused1, $unused2),
(TFStrings_PrintOp $value)>;
diff --git a/integrations/tensorflow/compiler/dialect/tf_strings/conversion/test/tf_strings_to_strings.mlir b/integrations/tensorflow/compiler/dialect/tf_strings/conversion/test/tf_strings_to_strings.mlir
index 5492142..a220ae3 100644
--- a/integrations/tensorflow/compiler/dialect/tf_strings/conversion/test/tf_strings_to_strings.mlir
+++ b/integrations/tensorflow/compiler/dialect/tf_strings/conversion/test/tf_strings_to_strings.mlir
@@ -35,3 +35,12 @@
// CHECK: return [[VAL0]]
return %0 : !tf_strings.string
}
+
+// CHECK-LABEL: @gather
+func @gather(%arg0: tensor<5x!tf_strings.string>, %arg1: tensor<3xi32>) -> tensor<3x!tf_strings.string> {
+ // CHECK-DAG: [[VAL0:%.+]] = "strings.gather"(%arg0, %arg1)
+ %0 = "tf_strings.gather"(%arg0, %arg1) : (tensor<5x!tf_strings.string>, tensor<3xi32>) -> tensor<3x!tf_strings.string>
+
+ // CHECK: return [[VAL0]]
+ return %0 : tensor<3x!tf_strings.string>
+}
diff --git a/integrations/tensorflow/compiler/dialect/tf_strings/conversion/test/tf_to_tf_strings.mlir b/integrations/tensorflow/compiler/dialect/tf_strings/conversion/test/tf_to_tf_strings.mlir
index dc63d9a..eb42d54 100644
--- a/integrations/tensorflow/compiler/dialect/tf_strings/conversion/test/tf_to_tf_strings.mlir
+++ b/integrations/tensorflow/compiler/dialect/tf_strings/conversion/test/tf_to_tf_strings.mlir
@@ -36,3 +36,11 @@
return
}
+// CHECK-LABEL: func @gatherv2
+func @gatherv2(%arg0: tensor<5x!tf.string>, %arg1: tensor<3xi32>) {
+ %0 = "tf.Const"() {device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+
+ // CHECK-DAG: [[VAL0:%.+]] = "tf_strings.gather"(%arg0, %arg1)
+ %1 = "tf.GatherV2"(%arg0, %arg1, %0) {batch_dims = 0 : i64, device = ""} : (tensor<5x!tf.string>, tensor<3xi32>, tensor<i32>) -> tensor<3x!tf.string>
+ return
+}
diff --git a/integrations/tensorflow/compiler/dialect/tf_strings/ir/ops.td b/integrations/tensorflow/compiler/dialect/tf_strings/ir/ops.td
index 9cf9379..b3faff5 100644
--- a/integrations/tensorflow/compiler/dialect/tf_strings/ir/ops.td
+++ b/integrations/tensorflow/compiler/dialect/tf_strings/ir/ops.td
@@ -44,6 +44,34 @@
OpBuilderDAG<(ins "Value":$value)>];
}
+def TFStrings_GatherOp : Op<TFStrings_Dialect, "gather"> {
+ let summary = "Gathers all the strings from by index.";
+ let description = [{
+ Gathers all the strings from a Tensor using the index of the last dimension.
+ }];
+
+ let arguments = (ins TFStrings_StringTensor:$dict,
+ TFStrings_ValueTensor:$indices);
+
+ let results = (outs
+ TFStrings_StringTensor:$result
+ );
+}
+
+def TFStrings_ConcatOp : Op<TFStrings_Dialect, "concat"> {
+ let summary = "Concatenates the strings in the tensor along the last dimension.";
+
+ let description = [{
+ Concatenates the strings in the tensor along the last dimension.
+ }];
+
+ let arguments = (ins TFStrings_StringTensor:$value);
+
+ let results = (outs
+ TFStrings_StringTensor:$result
+ );
+}
+
def TFStrings_ToStringTensorOp : TFStrings_Op<"to_string_tensor"> {
let summary = "Converts a tensor of values to a tensor of strings.";
let description = [{
diff --git a/integrations/tensorflow/e2e/strings_test.py b/integrations/tensorflow/e2e/strings_test.py
index 523062d..e8ef23d 100644
--- a/integrations/tensorflow/e2e/strings_test.py
+++ b/integrations/tensorflow/e2e/strings_test.py
@@ -33,6 +33,13 @@
string_tensor = tf.strings.as_string(ids)
tf.print(string_tensor)
+ @tf.function(input_signature=[
+ tf.TensorSpec((None,), dtype=tf.int32),
+ tf.TensorSpec((None,), dtype=tf.int32),
+ ])
+ def gather(self, string_values, indices):
+ tf.print(tf.gather(tf.as_string(string_values), indices))
+
# @tf.function(input_signature=[
# tf.TensorSpec((None, None), dtype=tf.int32),
# ])
@@ -56,6 +63,17 @@
self.compare_backends(print_ids, self._modules)
+ def test_gather(self):
+
+ def gather(module):
+ string_values = np.asarray([ord(c) for c in string.printable],
+ dtype=np.int32)
+ input_indices = np.asarray([12, 10, 29, 21, 10, 34], dtype=np.int32)
+ module.gather(string_values, input_indices)
+
+ self.compare_backends(gather, self._modules)
+
+
# def test_strings_to_ids(self):
#
# def strings_to_ids(module):
diff --git a/iree/compiler/Dialect/Modules/Strings/Conversion/ConversionPatterns.cc b/iree/compiler/Dialect/Modules/Strings/Conversion/ConversionPatterns.cc
index e02895c..557bc85 100644
--- a/iree/compiler/Dialect/Modules/Strings/Conversion/ConversionPatterns.cc
+++ b/iree/compiler/Dialect/Modules/Strings/Conversion/ConversionPatterns.cc
@@ -29,6 +29,9 @@
patterns.insert<HALOpConversion<IREE::Strings::ToStringTensorOp,
IREE::Strings::ToStringTensorOp>>(
context, typeConverter);
+ patterns.insert<
+ HALOpConversion<IREE::Strings::GatherOp, IREE::Strings::GatherOp>>(
+ context, typeConverter);
}
void populateStringsToVMPatterns(MLIRContext *context,
diff --git a/iree/compiler/Dialect/Modules/Strings/IR/Ops.td b/iree/compiler/Dialect/Modules/Strings/IR/Ops.td
index 743c351..92f8265 100644
--- a/iree/compiler/Dialect/Modules/Strings/IR/Ops.td
+++ b/iree/compiler/Dialect/Modules/Strings/IR/Ops.td
@@ -94,13 +94,13 @@
}
def STRINGS_GatherOp : Op<STRINGS_Dialect, "gather", [NoSideEffect]> {
- let summary = [{gathers all the strings from a Tensor by id}];
+ let summary = "Gathers all the strings from a Tensor by index.";
let description = [{
- Gathers all the strings from a Tensor by ID.
+ Gathers all the strings from a Tensor by index values along the final axis.
}];
let arguments = (ins STRINGS_StringTensor:$dict,
- STRINGS_ValueTensor:$ids);
+ STRINGS_ValueTensor:$indices);
let results = (outs
STRINGS_StringTensor:$result
@@ -108,9 +108,7 @@
}
def STRINGS_ConcatOp : Op<STRINGS_Dialect, "concat", [NoSideEffect]> {
- let summary = [{
- concatenates the strings in the tensor along the last dimension
- }];
+ let summary = "Concatenates a tensor of strings along the last dimension";
let description = [{
Concatenates the strings in the tensor along the last dimension.
@@ -124,7 +122,7 @@
}
def STRINGS_PrintOp : Op<STRINGS_Dialect, "print"> {
- let summary = [{prints the contents of a string}];
+ let summary = "Prints the contents of a string.";
let description = [{
Prints the contents of a string.
}];
