[llvm-branch-commits] [mlir] 7a34031 - [MLIR] Add verifier for LHLO get tuple element op

[Uday Bondhugula via llvm-branch-commits]

Author: Uday Bondhugula
Date: 2021-09-23T05:52:22+05:30
New Revision: 7a34031a8fdabe5f484a537c96a7cd31999c51d5

URL: https://github.com/llvm/llvm-project/commit/7a34031a8fdabe5f484a537c96a7cd31999c51d5
DIFF: https://github.com/llvm/llvm-project/commit/7a34031a8fdabe5f484a537c96a7cd31999c51d5.diff

LOG: [MLIR] Add verifier for LHLO get tuple element op

The LHLO GetTupleElementOp was missing its verifier. Add the verifier.

Added: 
    mlir/test/Dialect/LHLO/invalid.mlir

Modified: 
    mlir/include/mlir/Dialect/LHLO/IR/LHLOOps.td
    mlir/lib/Dialect/LHLO/IR/LHLOOps.cc

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/LHLO/IR/LHLOOps.td b/mlir/include/mlir/Dialect/LHLO/IR/LHLOOps.td
index 15c17ee3b85c..104e9c011de6 100644
--- a/mlir/include/mlir/Dialect/LHLO/IR/LHLOOps.td
+++ b/mlir/include/mlir/Dialect/LHLO/IR/LHLOOps.td
@@ -251,6 +251,8 @@ def LHLO_GetTupleElementOp: LHLO_Op<"get_tuple_element", []>, BASE_HLO_GetTupleE
   );
 
   let results = (outs Arg<LHLO_BufferOrIntOrFP>);
+
+  let verifier = [{ return ::verify(*this); }];
 }
 
 def LHLO_CompareOp: LHLO_Op<"compare", []>, BASE_HLO_CompareOp {

diff  --git a/mlir/lib/Dialect/LHLO/IR/LHLOOps.cc b/mlir/lib/Dialect/LHLO/IR/LHLOOps.cc
index 26251379be8f..01b0da948dea 100644
--- a/mlir/lib/Dialect/LHLO/IR/LHLOOps.cc
+++ b/mlir/lib/Dialect/LHLO/IR/LHLOOps.cc
@@ -59,11 +59,10 @@ void LHLODialect::initialize() {
       >();
 }
 
-#define GET_OP_CLASSES
-#include "mlir/Dialect/LHLO/IR/LHLOOps.cpp.inc"
-
-namespace mlir {
-namespace xla_lhlo {
+using xla_lhlo::ConstOp;
+using xla_lhlo::FusionOp;
+using xla_lhlo::GetTupleElementOp;
+using xla_lhlo::TupleOp;
 
 //===----------------------------------------------------------------------===//
 // ConstOp.
@@ -123,5 +122,30 @@ void TupleOp::build(OpBuilder& builder, OperationState& result,
   build(builder, result, builder.getTupleType(types), values);
 }
 
-}  // namespace xla_lhlo
-}  // namespace mlir
+//===----------------------------------------------------------------------===//
+// GetTupleElementOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(GetTupleElementOp op) {
+  unsigned index = op.index();
+  auto operandType = op.getOperand().getType().cast<TupleType>();
+  if (index >= operandType.size()) {
+    return op.emitOpError(
+        llvm::formatv("index {0} is out of bounds of operand with size {1}",
+                      index, operandType.size()));
+  }
+
+  Type expectedType = operandType.getType(index);
+  if (op.getType() != expectedType) {
+    return op.emitOpError(llvm::formatv("has return type {0}, but expected {1}",
+                                        op.getType(), expectedType));
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/LHLO/IR/LHLOOps.cpp.inc"

diff  --git a/mlir/test/Dialect/LHLO/invalid.mlir b/mlir/test/Dialect/LHLO/invalid.mlir
new file mode 100644
index 000000000000..0e269535d00c
--- /dev/null
+++ b/mlir/test/Dialect/LHLO/invalid.mlir
@@ -0,0 +1,10 @@
+// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+
+func @passthrough(%arg : memref<8xi32>) {
+  %c0_i32 = constant 0 : i32
+  %tuple = "xla_lhlo.tuple"(%c0_i32, %arg) : (i32, memref<8xi32>) -> (tuple<i32, memref<8xi32>>)
+  %elt = "xla_lhlo.get_tuple_element"(%tuple) {index = 0 : i32} : (tuple<i32, memref<8xi32>>) -> i32
+  // expected-error at +1{{'xla_lhlo.get_tuple_element' op has return type memref<8xi32>, but expected i32}}
+  %mem = "xla_lhlo.get_tuple_element"(%tuple) {index = 0 : i32} : (tuple<i32, memref<8xi32>>) -> memref<8xi32>
+  return
+}