Add free standing getElementTypeOrSelf member.

This function returns the element type of the underlying type or the input type itself. This removes some of the casting and code from ODS and into C++ code.

I've not converted all the call sites (as this requires a new include and target) and wanted to run it past folks first.

PiperOrigin-RevId: 251978156

diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index dd4c276..a170069 100644 (file)
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -1024,8 +1024,7 @@ class TCopVTEtIs<int idx, Type type> : And<[
 // given type.
 class ArgOrResultElementTypeIs<string name, Type type> : And<[
    SubstLeaves<"$_self",  "$" # name # ".getType()", IsShapedTypePred>,
-   SubstLeaves<"$_self",  "$" # name #
-     ".getType().cast<ShapedType>().getElementType()",
+   SubstLeaves<"$_self",  "mlir::getElementTypeOrSelf($" # name # ")",
      type.predicate>]>;
 
 // Predicate to verify that the i'th operand and the j'th operand have the same
@@ -1033,15 +1032,13 @@ class ArgOrResultElementTypeIs<string name, Type type> : And<[
 // Type Constraint operand `i`'s Element type is Same As operand `j`'s Element
 // type.
 class TCopVTEtIsSameAs<int i, int j> : And<[
-    CPred<"$_op.getNumOperands() > std::max(" # i # "," # j # ")">,
+    CPred<"$_op.getNumOperands() > std::max(" # i # "u," # j # "u)">,
     SubstLeaves<"$_self",  "$_op.getOperand(" # i # ")->getType()",
       IsShapedTypePred>,
     SubstLeaves<"$_self",  "$_op.getOperand(" # j # ")->getType()",
       IsShapedTypePred>,
-    // TODO: This could be made into C++ function instead.
-    CPred<"$_op.getOperand(" # i # ")->getType().cast<ShapedType>()."
-          "getElementType() == $_op.getOperand(" # j # ")->getType()."
-          "cast<ShapedType>().getElementType()">]>;
+    CPred<"mlir::getElementTypeOrSelf($_op.getOperand(" # i # ")) == "
+          "mlir::getElementTypeOrSelf($_op.getOperand(" # j # "))">]>;
 
 // Predicate to verify that the i'th result and the j'th operand have the same
 // elemental type.

diff --git a/mlir/include/mlir/Support/TypeUtilities.h b/mlir/include/mlir/Support/TypeUtilities.h
new file mode 100644 (file)
index 0000000..255281f
--- /dev/null
+++ b/mlir/include/mlir/Support/TypeUtilities.h
@@ -0,0 +1,41 @@
+//===- TypeUtilities.h - Helper function for type queries -------*- C++ -*-===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This file defines generic type utilities.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_SUPPORT_TYPEUTILITIES_H
+#define MLIR_SUPPORT_TYPEUTILITIES_H
+
+namespace mlir {
+
+class Attribute;
+class Type;
+class Value;
+
+/// Return the element type or return the type itself.
+Type getElementTypeOrSelf(Type type);
+
+/// Return the element type or return the type itself.
+Type getElementTypeOrSelf(Attribute attr);
+Type getElementTypeOrSelf(Value *val);
+Type getElementTypeOrSelf(Value &val);
+
+} // end namespace mlir
+
+#endif // MLIR_SUPPORT_TYPEUTILITIES_H

diff --git a/mlir/lib/Support/CMakeLists.txt b/mlir/lib/Support/CMakeLists.txt
index 97da45b..bac4854 100644 (file)
--- a/mlir/lib/Support/CMakeLists.txt
+++ b/mlir/lib/Support/CMakeLists.txt
@@ -1,3 +1,9 @@
+set(LLVM_OPTIONAL_SOURCES
+  FileUtilities.cpp
+  StorageUniquer.cpp
+  TypeUtilities.cpp
+)
+
 add_llvm_library(MLIRSupport
   FileUtilities.cpp
   StorageUniquer.cpp
@@ -6,3 +12,11 @@ add_llvm_library(MLIRSupport
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Support
   )
 target_link_libraries(MLIRSupport LLVMSupport)
+
+add_llvm_library(MLIRTypeUtilities
+  TypeUtilities.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Support
+  )
+target_link_libraries(MLIRTypeUtilities MLIRIR MLIRStandardOps)

diff --git a/mlir/lib/Support/TypeUtilities.cpp b/mlir/lib/Support/TypeUtilities.cpp
new file mode 100644 (file)
index 0000000..8114b73
--- /dev/null
+++ b/mlir/lib/Support/TypeUtilities.cpp
@@ -0,0 +1,46 @@
+//===- TypeUtilities.h - Helper function for type queries -----------------===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This file defines generic type utilities.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Support/TypeUtilities.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/StandardTypes.h"
+#include "mlir/IR/Types.h"
+#include "mlir/IR/Value.h"
+
+using namespace mlir;
+
+Type mlir::getElementTypeOrSelf(Type type) {
+  if (auto st = type.dyn_cast<ShapedType>())
+    return st.getElementType();
+  return type;
+}
+
+Type mlir::getElementTypeOrSelf(Value *val) {
+  return getElementTypeOrSelf(val->getType());
+}
+
+Type mlir::getElementTypeOrSelf(Value &val) {
+  return getElementTypeOrSelf(val.getType());
+}
+
+Type mlir::getElementTypeOrSelf(Attribute attr) {
+  return getElementTypeOrSelf(attr.getType());
+}

diff --git a/mlir/test/TestDialect/CMakeLists.txt b/mlir/test/TestDialect/CMakeLists.txt
index c08b74e..59fc920 100644 (file)
--- a/mlir/test/TestDialect/CMakeLists.txt
+++ b/mlir/test/TestDialect/CMakeLists.txt
@@ -33,5 +33,6 @@ whole_archive_link(mlir-test-opt
 target_link_libraries(mlir-test-opt
   PRIVATE
   MLIRMlirOptLib
+  MLIRTypeUtilities
   LLVMSupport
 )

diff --git a/mlir/test/TestDialect/TestDialect.cpp b/mlir/test/TestDialect/TestDialect.cpp
index 49966df..d91637c 100644 (file)
--- a/mlir/test/TestDialect/TestDialect.cpp
+++ b/mlir/test/TestDialect/TestDialect.cpp
@@ -17,6 +17,7 @@
 
 #include "TestDialect.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/Support/TypeUtilities.h"
 
 using namespace mlir;
 

diff --git a/mlir/test/TestDialect/TestOps.td b/mlir/test/TestDialect/TestOps.td
index b0296e1..845b08d 100644 (file)
--- a/mlir/test/TestDialect/TestOps.td
+++ b/mlir/test/TestDialect/TestOps.td
@@ -96,6 +96,12 @@ def ArgAndResHaveFixedElementTypesOp :
   let results = (outs AnyVectorOrTensor:$res);
 }
 
+def OperandsHaveSameElementType : TEST_Op<"operands_have_same_element_type", [
+    PredOpTrait<"first and second operand have same element type",
+                TCopVTEtIsSameAs<0, 1>>]> {
+  let arguments = (ins AnyTensor:$x, AnyTensor:$y);
+}
+
 //===----------------------------------------------------------------------===//
 // Test Patterns
 //===----------------------------------------------------------------------===//

diff --git a/mlir/test/mlir-tblgen/predicate.td b/mlir/test/mlir-tblgen/predicate.td
index 9c37e93..454a01b 100644 (file)
--- a/mlir/test/mlir-tblgen/predicate.td
+++ b/mlir/test/mlir-tblgen/predicate.td
@@ -35,19 +35,6 @@ def OpC : NS_Op<"op_for_TCopVTEtIs", [
 // CHECK-LABEL: OpC::verify
 // CHECK: if (!((((*this->getOperation()).getNumOperands() > 0)) && (((*this->getOperation()).getOperand(0)->getType().isa<ShapedType>())) && (((*this->getOperation()).getOperand(0)->getType().cast<ShapedType>().getElementType().isInteger(32)))))
 
-
-def OpD : NS_Op<"op_for_TCOpVTEtIsSameAs", [
-    PredOpTrait<"first operand is a shaped type with the same "
-                "elemental type as itself",
-                TCopVTEtIsSameAs<0, 0>>]> {
-  let arguments = (ins AnyTensor:$x);
-}
-
-// CHECK-LABEL: OpD::verify
-// CHECK: if (!((((*this->getOperation()).getNumOperands() > std::max(0,0))) && (((*this->getOperation()).getOperand(0)->getType().isa<ShapedType>())) && (((*this->getOperation()).getOperand(0)->getType().isa<ShapedType>())) && (((*this->getOperation()).getOperand(0)->getType().cast<ShapedType>().getElementType() == (*this->getOperation()).getOperand(0)->getType().cast<ShapedType>().getElementType()))))
-// CHECK-NEXT: return emitOpError("failed to verify that first operand is a shaped type with the same elemental type as itself");
-
-
 def OpE : NS_Op<"op_for_TCresVTEtIsSameAsOp", [
     PredOpTrait<"first operand is a shaped type with the same "
                 "elemental type as first result",

diff --git a/mlir/test/mlir-tblgen/types.mlir b/mlir/test/mlir-tblgen/types.mlir
index e84adcf..d306365 100644 (file)
--- a/mlir/test/mlir-tblgen/types.mlir
+++ b/mlir/test/mlir-tblgen/types.mlir
@@ -110,3 +110,11 @@ func @fixed_element_types(%arg0: tensor<* x i32>, %arg1: tensor<* x f32>) {
   %0 = "test.arg_and_res_have_fixed_element_types"(%arg1, %arg0) {attr: ""}: (tensor<* x f32>, tensor<* x i32>) -> tensor<* x i16>
   return
 }
+
+// -----
+
+func @fixed_element_types(%arg0: tensor<* x i32>, %arg1: tensor<* x f32>) {
+  // expected-error@+1 {{failed to verify that first and second operand have same element type}}
+  "test.operands_have_same_element_type"(%arg1, %arg0): (tensor<* x f32>, tensor<* x i32>) -> ()
+  return
+}