Compilation failure with DRR-generated pattern

MLIR
1

I was trying to write the following declarative rewrite rule (tensor_ext.rotate is my custom out-of-tree op):

def GetZeroAttr : NativeCodeCall<"$_builder.getIndexAttr(0)">;

def InsertRotations : Pattern<
  (Arith_AddIOp
    (Tensor_ExtractOp $t1, (variadic $i1)),
    (Tensor_ExtractOp $t2, (variadic $i2)),
    $overflow),
  [
    (Arith_ConstantOp:$zero (GetZeroAttr)),
    (TensorExt_RotateOp:$r1 $t1, $i1),
    (TensorExt_RotateOp:$r2 $t2, $i2),
    (Arith_AddIOp:$addResult $r1, $r2, $overflow),
    (Tensor_ExtractOp $addResult, $zero),
  ]
>;

And the generated code raises the following error during compilation:

error: invalid range expression of type 'mlir::OpResult'; no viable 'begin' function available
  123 |       for (auto v: (*zero.getODSResults(0).begin())) {
      |                  ^ ~

More specifically, it seems to be trying to convert from mlir::Value to Attribute or NamedAttribute when it tries to iterate. The snippet it complains about:

    ::mlir::tensor::ExtractOp tblgen_ExtractOp_1;
    {
      ::llvm::SmallVector<::mlir::Value, 4> tblgen_values; (void)tblgen_values;
      ::llvm::SmallVector<::mlir::NamedAttribute, 4> tblgen_attrs; (void)tblgen_attrs;
      tblgen_values.push_back((*addResult.getODSResults(0).begin()));
      for (auto v: (*zero.getODSResults(0).begin())) {
        tblgen_values.push_back(v);
      }
      ::llvm::SmallVector<::mlir::Type, 4> tblgen_types; (void)tblgen_types;
      for (auto v: castedOp0.getODSResults(0)) {
        tblgen_types.push_back(v.getType());
      }
      tblgen_ExtractOp_1 = rewriter.create<::mlir::tensor::ExtractOp>(odsLoc, tblgen_types, tblgen_values, tblgen_attrs);
    }

I’m confused why this is generated this way. Is there something I’ve done wrong in setting up the pattern? Other attempts to phrase the constant op (e.g., ConstantAttr<IndexAttr, "0">) also fail with this error.

Here is ths full error trace: gist:68ac2099d9e6d4a983e614bf1521b2d3 · GitHub

2

I get the same problem with the following pattern that makes me think it is probably a bug with how variadic operands are handled.

def RotatePlusExtractToIndexedExtract : Pat<
  (Tensor_ExtractOp
    (TensorExt_RotateOp $tensor, $shift),
    (variadic $index)),
  (Tensor_ExtractOp
    $tensor,
    (Arith_AddIOp $shift, $index, DefOverflow))
>;
3

Yes, variadic operands are not supported today (there should have been an error …), you could do it “manually”:

def GetZeroIndexAttr : NativeCodeCall<"$_builder.getIndexAttr(0)">;
def MakeSingleResultVariadic: NativeCodeCall<"{ $0 }">;

def InsertRotations : Pattern<
  (Arith_AddIOp
    (Tensor_ExtractOp $t1, (variadic $i1)),
    (Tensor_ExtractOp $t2, (variadic $i2)),
    $overflow),
  [
    (Arith_ConstantOp:$zero (GetZeroIndexAttr)),
    (Arith_AddIOp:$addResult $t1, $t2, $overflow),
    (Tensor_ExtractOp $addResult,
        (MakeSingleResultVariadic $zero)),
  ]
>;

When variadic was added for matching, it was not added to result side too. So two improvements desirable here:

  1. Error message during generation.
  2. variadic keyword for result side.
1 Like
4

Works beautifully! Let’s chat more later about how to improve this and I can do it.