[mlir] Minor fixes after removing types from attributes

[dcaballe]

Summary

D130092 removed types from attributes. This patch fixes a minor
issues what was exposed when integrating that change in IREE. An
explicit cast is needed so that the template type of makeArrayRef
is automatically deduced.

[dcaballe]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1249 ↗

(On Diff #451570)

Hey @ThomasRaoux, this AllTypesMatch here is not looking into the variadic container and comparing the types in there. It's just comparing the type of the container. Could you please follow up with a fix for this? I couldn't find an existing predicate to do that.

[dcaballe]

Actually, this is not working when we actually compare the type of multiple variadic operands with each other. For example:

 def SimdLoopOp : OpenMP_Op<"simdloop", [AttrSizedOperandSegments,
                        AllTypesMatch<["lowerBound", "upperBound", "step"]>]> {
let summary = "simd loop construct";
 let description = [{...}];

 // TODO: Add other clauses
 let arguments = (ins Variadic<IntLikeType>:$lowerBound,
            Variadic<IntLikeType>:$upperBound,
            Variadic<IntLikeType>:$step,
            Optional<I1>:$if_expr,
            Confined<OptionalAttr<I64Attr>, [IntPositive]>:$simdlen,
            UnitAttr:$inclusive
    );

I would appreciate some help here.
Thanks!

[dcaballe]

To provide more context (I talked to Jeff privately), D130092 introduces some changes where getType returns a ShapedType instead of Type. This leads to some problem in ODS, when we use SameTypeOperands or AllTypeMatch and the getType() for the values to compare result in a mix of ShapedType and Type: we use llvm::makeArrayRef(...) and the template type can't be deduced automatically.

In the following example, $lhs.getType() will return Type and $rhs.getType() will return ShapedType:

def CHECK_ExpectEqConstOp :
     Op<CHECK_Dialect, "expect_eq_const", [SameTypeOperands]> {
    let summary = [{Checks that the tensor operand is equal to some constant}];
    let description =  "";

    let arguments = (ins
      AnyTensor:$lhs,
      ElementsAttr:$value
    );

    let hasCanonicalizer = 1;

    let assemblyFormat = "`(` $lhs `,` $value `)` attr-dict `:` type($lhs)";
  }

Then, the generated llvm::makeArrayRef won't compile:

if (!((::llvm::is_splat(::llvm::makeArrayRef({(*this->getODSOperands(0).begin()).getType(), value().getType()})))))
  return emitOpError("failed to verify that all of {lhs, value} have same type");           
return ::mlir::success();

note: candidate template ignored: deduced conflicting types for param
eter 'T' ('mlir::Type' vs. 'mlir::ShapedType')

Perhaps we should revert the changes related to ShapedType getType() in D130092?

@jpienaar @mehdi_amini

[Mogball]

Going forward, there will be no guarantees that getType will need to return Type. FloatAttr::getType should be changed to return FloatType, for example. The whole point is that an attribute with a type can return a more refined subclass if needed.

[Mogball]

mlir/include/mlir/IR/OpBase.td

2327

In the absolute worst case, you could do

makeArrayRef<std::conditional_t<std::is_base_of<::mlir::Type, decltype(...)>, Type, decltype(...)>>({...})

2327

But you could just introduce new traits for variadic and non-variadic types. Unfortunately, these traits have some interactions with OpDefinitionsGen.cpp as well

[jpienaar]

mlir/include/mlir/IR/OpBase.td

2327

Here we just want to compare the types correct, so casting to Type (in future getType could return many different types, but all Types yes?) seems to enable the comparison that we want here.

[dcaballe]

Use Jeff's suggestion.

[dcaballe]

Thanks for the suggestions! The last diff should work.

mlir/include/mlir/IR/OpBase.td

2327

That's why I initially thought but we are also using this trait to compare variadic arguments. getType() returns the type of the variadic container (ValueTypeRange?) so we can't convert everything to Type.

[Mogball]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1250 ↗

(On Diff #451614)

unrelated diff

mlir/include/mlir/IR/OpBase.td

2292

This works, but I don't think it's the best solution (even though I suggested it myself). Why don't we move away from llvm::is_splat and just expand to values[0] == values[1] && values[1] == values[2] etc. This way we don't have to rely on the fickle template parameter inference of makeArrayRef and initializer lists.

[rriddle]

mlir/include/mlir/IR/OpBase.td

2292

Yeah we should either expand, or cast all of the values to Type; i.e. by wrapping them with ::mlir::Type().

[dcaballe]

Expanded list values
Fix provided by @antiagainst.

[Mogball]

mlir/include/mlir/IR/OpBase.td

2294

Can you add braces around the lhs and rhs arguments?

[dcaballe]

Addressed feedback

[Mogball]

mlir/include/mlir/IR/OpBase.td

2290–2295

Does this work?

[dcaballe]

mlir/include/mlir/IR/OpBase.td

2290–2295

Unfortunately, it's not. The first == would return a bool that would be compared against a value and so on...

[Mogball]

LGTM

mlir/include/mlir/IR/OpBase.td

2290–2295

Right thanks

mlir/include/mlir/IR/OpBase.td

class Shape<string name> :
     StrFunc<"$" # name # ".getType().cast<::mlir::ShapedType>().getShape()">;
 
 class ElementCount<string name> :
   StrFunc<"$" # name # ".getType().cast<::mlir::ShapedType>()"
                                  ".getNumElements()">;
 
 class ElementType<string name> : StrFunc<"getElementTypeOrSelf($" # name # ")">;
 
 class AllMatchPred<list<string> values> :
-    CPred<"::llvm::is_splat(::llvm::makeArrayRef({"
-          # !interleave(values, ", ") #"}))">;
+  CPred<!if(!lt(!size(values), 2),
+            "true",
+            !foldl("(" # !head(values) # ")", !tail(values), acc, v,
+                   acc # " == (" # v # ") && (" # v # ")")
+              # " == (" # !head(values) # ")")>;
 
 class AllMatch<list<string> values, string summary> :
     PredOpTrait<summary, AllMatchPred<values>>;
 
 // TODO: Only works for non-variadic.
 class AllMatchSameOperatorPred<list<string> names, string operator> :
     AllMatchPred<!foreach(n, names, !subst("$_self", "$" # n, operator))>;
 

 
 class AllRanksMatch<list<string> names> :
     AllMatchSameOperatorTrait<names, Rank<"_self">.result, "rank">;
 
 class AllShapesMatch<list<string> names> :
     AllMatchSameOperatorTrait<names, Shape<"_self">.result, "shape">;
 
 class AllTypesMatch<list<string> names> :
     AllMatchSameOperatorTrait<names, "$_self.getType()", "type">;
 
 // A type constraint that denotes `transform(lhs.getType()) == rhs.getType()`.
 // An optional comparator function may be provided that changes the above form
 // into: `comparator(transform(lhs.getType()), rhs.getType())`.
 class TypesMatchWith<string summary, string lhsArg, string rhsArg,
                      string transform, string comparator = "std::equal_to<>()">
   : PredOpTrait<summary, CPred<
       comparator # "(" #