[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

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

2316

In the absolute worst case, you could do

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

2316

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

2316

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

2316

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

unrelated diff

mlir/include/mlir/IR/OpBase.td

2283

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

2283

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

2287

Can you add braces around the lhs and rhs arguments?

[dcaballe]

Addressed feedback

[Mogball]

mlir/include/mlir/IR/OpBase.td

2280–2285

Does this work?

[dcaballe]

mlir/include/mlir/IR/OpBase.td

2280–2285

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

2280–2285

Right thanks

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

def MMAElementWise : I32EnumAttr<"MMAElementwiseOp",
   let genSpecializedAttr = 0;
   let cppNamespace = "::mlir::gpu";
 }
 def MMAElementWiseAttr : EnumAttr<GPU_Dialect, MMAElementWise,
                                   "mma_element_wise">;
 
 def GPU_SubgroupMmaElementwiseOp : GPU_Op<"subgroup_mma_elementwise",
     [NoSideEffect,
      AllTypesMatch<["args"]>]>{
-
   let summary = "GPU warp elementwise operation on a matrix";
 
   let description = [{
     The `gpu.subgroup_mma_elementwise` takes `!gpu.mma_matrix` inputs and
     compute a new `!gpu.mma_matrix` by applying an elementwise operation to each
     element.
 
     Since the operation is elementwise and the matrix type must match, the

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({"
+    CPred<"::llvm::is_splat(::llvm::makeArrayRef<std::conditional_t<"
+            # "std::is_base_of<::mlir::Type, decltype(" # !head(values) # ")>::value,"
+            # "::mlir::Type, decltype(" # !head(values) # ")>>({"
           # !interleave(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 AllMatchSameOperatorTrait<list<string> names, string operator,
                                 string summary> :
     PredOpTrait<

 
 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 # "(" #