[mlir][ods] Use lambda in element type check pred rather than repeated casts

[jpienaar]

Summary

Avoids multiple cast & getElementType calls. Just a local change for ShapedType
containers but reduces one run from 24.7 to 24.04s.

Resultant code generated change:
https://gist.github.com/jpienaar/7ffd2e9b0737134ba2ea2729b91c9572

[rriddle]

mlir/include/mlir/IR/OpBase.td

570–574

Was this used by anything else?

575

mlir/include/mlir/IR/OpBase.td

 // A container type is a type that has another type embedded within it.
 class ContainerType<Type etype, Pred containerPred, code elementTypeCall,
                     string descr, string cppClassName = "::mlir::Type"> :
     // First, check the container predicate.  Then, substitute the extracted
     // element into the element type checker.
     Type<And<[containerPred,
                 SubstLeaves<"$_self", !cast<string>(elementTypeCall),
                 etype.predicate>]>,
-         descr # " of " # etype.summary # " values", cppClassName> {
-  // The type of elements in the container.
-  Type elementType = etype;
-
-  // Call to retrieve.
-  code getElementTypeCall = elementTypeCall;
-}
+         descr # " of " # etype.summary # " values", cppClassName>;
 
 class ShapedContainerType<list<Type> allowedTypes,
                           Pred containerPred, string descr,
                           string cppClassName = "::mlir::Type"> :
-    ContainerType<AnyTypeOf<allowedTypes>, containerPred,
-                  "$_self.cast<::mlir::ShapedType>().getElementType()", descr,
-                  cppClassName>;
+    Type<And<[containerPred,
+              Concat<"[](Type elementType) { return ",
+                SubstLeaves<"$_self", "elementType",
+                AnyTypeOf<allowedTypes>.predicate>,
+                "; }($_self.cast<::mlir::ShapedType>().getElementType())">]>,
+         descr # " of " # AnyTypeOf<allowedTypes>.summary # " values", cppClassName>;
 
 // Whether a shaped type is ranked.
 def HasRankPred : CPred<"$_self.cast<::mlir::ShapedType>().hasRank()">;
 
 // Whether a shaped type has one of the specified ranks.
 class HasAnyRankOfPred<list<int> ranks> : And<[
     HasRankPred,
     Or<!foreach(rank, ranks,

mlir/test/mlir-tblgen/predicate.td

 // CHECK-NEXT:  if (!((type.isInteger(32) || type.isF32()))) {
 // CHECK-NEXT:    return op->emitOpError(valueKind) << " #" << valueGroupStartIndex << " must be 32-bit integer or floating-point type, but got " << type;
 
 // Check there is no verifier with same predicate generated.
 // CHECK-NOT:  if (!((type.isInteger(32) || type.isF32()))) {
 // CHECK-NOT:    return op->emitOpError(valueKind) << " #" << valueGroupStartIndex << " must be 32-bit integer or floating-point type, but got " << type;
 
 // CHECK: static ::mlir::LogicalResult [[$TENSOR_CONSTRAINT:__mlir_ods_local_type_constraint.*]](
-// CHECK-NEXT:  if (!(((type.isa<::mlir::TensorType>())) && ((true)))) {
+// CHECK-NEXT:  if (!(((type.isa<::mlir::TensorType>())) && ([](Type elementType) { return (true); }(type.cast<::mlir::ShapedType>().getElementType())))) {
 // CHECK-NEXT:    return op->emitOpError(valueKind) << " #" << valueGroupStartIndex << " must be tensor of any type values, but got " << type;
 
 // CHECK: static ::mlir::LogicalResult [[$TENSOR_INTEGER_FLOAT_CONSTRAINT:__mlir_ods_local_type_constraint.*]](
-// CHECK-NEXT:  if (!(((type.isa<::mlir::TensorType>())) && (((type.cast<::mlir::ShapedType>().getElementType().isF32())) || ((type.cast<::mlir::ShapedType>().getElementType().isSignlessInteger(32)))))) {
+// CHECK-NEXT:  if (!(((type.isa<::mlir::TensorType>())) && ([](Type elementType) { return ((elementType.isF32())) || ((elementType.isSignlessInteger(32))); }(type.cast<::mlir::ShapedType>().getElementType())))) {
 // CHECK-NEXT:    return op->emitOpError(valueKind) << " #" << valueGroupStartIndex << " must be tensor of 32-bit float or 32-bit signless integer values, but got " << type;
 
 // CHECK-LABEL: OpA::verify
 // CHECK: auto valueGroup0 = getODSOperands(0);
 // CHECK: for (::mlir::Value v : valueGroup0) {
 // CHECK:   if (::mlir::failed([[$INTEGER_FLOAT_CONSTRAINT]]
 
 def OpB : NS_Op<"op_for_And_PredOpTrait", [