Replacing one op with multiple ops

MLIR
1

Hi

I have a question w.r.t. 1-N mapping where converting between two dialects.

Specifically, when one op has to result in multiple ops and the results of those multiple ops have to replace the result of that one op.

I am using this current approach within matchAndRewrite: but i can only do 1-1 mappings…

rewriter.replaceOp(origO, {convertedOp.getODSResults(0)});
  • Is it possible to gather all results for those ops and use them in the replaceOp? - im currently getting this:
Assertion failed: (newValues.size() == op->getNumResults())
1 Like
2

Yes, it is possible to replace an op with a list of values. replaceOp takes a ValueRange as a second argument, which can be constructed from, e.g., a vector. As the assertion indicates, the list must contain as many values as the replaced op has results.

Also note that you shouldn’t normally use getODSResults in your code, it is an internal implementation detail.

3

Thanks. How can i replace an op with one results with multiple ops with multiple results? - do i need a higher order op that acts as a grouping op?

Also note that you shouldn’t normally use getODSResults in your code, it is an internal implementation detail.

How would you recommend passing the result?

4

replaceOp in fact replaces all uses of results (which are values) of an operation with other values provided as its argument. It doesn’t matter where do the values come from: one op, different ops, block arguments.

If you gave them names, ODS will produce named accessors. Same as for arguments.

5

Thanks a lot for your responses.

I still have an issue with getting something like this to compile:

    std::vector<MyType> convertedOps;
    for (auto ....) {
      auto op = rewriter.create<MyType>(...);
      convertedOps.emplace_back(op);
   }
    // original has only one result value
    rewriter.replaceOp(original, // how to use the convertedOps here? );
    return success();
6

It’s unclear what IR you want to produce.

rewriter.replaceOp(original, convertedOps.front().getResult(0)) will type check, but I have no idea if this is something you want to do or not.

7

Basically i have the following code:

%a = foo.bar {id = 0 : i32} : !foo.bar

which will get converted to:

%a1 = goo.bar {id = 0 : i32} : !goo.bar
%a2 = goo.bar {id = 1 : i32} : !goo.bar
%a3 = goo.bar {id = 2 : i32} : !goo.bar

Now i want to convert this:

foo.baz(%a) : (!foo.bar) -> ()

To:

goo.baz(%a) : (xxx) -> ()

How can i make xxx represent all the three goo.bars created in the second block? - do i need some higher-level representation that represents the whole group of goo.bars?

8

foo.bar returns a single value, the framework expects you to provide a replacement for exactly one value.

You’re using %a after your replacement here, but before you created %a1, %a2 and %a3: which one do you want to use?

9

Thanks. this is what i was trying to make sure of. so there is only 1-1 mapping.

I would like to access all of them as operands. What other approach could i take in order to represent a collection of types as an input? can i define my type in a certain way to represent a collection of types?

10

Can you provide the IR before and after and describe the conversion you’re trying to achieve?
Also type and operation are somehow orthogonal here: you can convert the operation foo.baz to goo.baz without seeing any type change (I think it’d help if you avoid using the same name for the operations and the types in your code sample).

11

Sure. here it is:

%a = foo.create_core {id = 0 : i32} : !foo.core // create an instance of core

Converted to:

%b1 = bar.create_thread {id = 0 : i32} : !bar.thread // create an instance of thread
%b2 = bar.create_thread {id = 1 : i32} : !bar.thread // create an instance of thread
%b3 = bar.create_thread {id = 2 : i32} : !bar.thread // create an instance of thread

Now this is the issue:

foo.execute (%a) : (!foo.core) -> ()

Converted: - this currently is not working

bar.schedule (%b1, %b2, %b3) : (!bar.thread,  !bar.thread,  !bar.thread) -> ()
12

You cannot replace one value with many values, at least not currently. This can be done in a staged way with a dummy op.

First,

%a = foo.create_core {id = 0 : i32} : !foo.core // create an instance of core
foo.execute (%a) : (!foo.core) -> ()

gets converted into

%b1 = bar.create_thread {id = 0 : i32} : !bar.thread // create an instance of thread
%b2 = bar.create_thread {id = 1 : i32} : !bar.thread // create an instance of thread
%b3 = bar.create_thread {id = 2 : i32} : !bar.thread // create an instance of thread
%b = dummy.op(%b1, %b2, %b3) : tuple<!bar.thread, !bar.thread, !bar.thread>
foo.execute (%b) : (tuple<!bar.thread, !bar.thread, !bar.thread>) -> ()

then the pattern matching foo.execute can check if the new operand given to it is defined by the dummy op and, if so, take the operands of that op instead, before erasing the dummy op. Since the dummy op was created during the conversion, it’s should be okay to use its operands.

13

Thanks. this is very helpful.

14

I tried this with my own type: but its throwing errors:

here is what i have:


struct ThreadGroupTypeStorage : public TypeStorage {
  ThreadGroupTypeStorage(ArrayRef<Thread> threads) : threads(threads) {}

  using KeyTy = ArrayRef<Thread>;

  static auto hashKey(const KeyTy &key) -> llvm::hash_code {
    return llvm::hash_value(key);
  }

  auto operator==(const KeyTy &key) const -> bool {
    return KeyTy(threads) == key;
  }

  static auto construct(TypeStorageAllocator &allocator, KeyTy key)
      -> ThreadGroupTypeStorage * {
    return new (allocator.allocate<ThreadGroupTypeStorage>())
        ThreadGroupTypeStorage(key);
  }

  ArrayRef<Thread> threads;
};

struct ThreadTypeStorage : public TypeStorage {
  ThreadTypeStorage(ThreadType type) : type(type) {}

  /// The hash key used for uniquing, just type
  using KeyTy = ThreadType;

  static auto hashKey(const KeyTy &key) -> llvm::hash_code {
    return llvm::hash_value(key);
  }

  auto operator==(const KeyTy &key) const -> bool { return KeyTy(type) == key; }

  static auto construct(TypeStorageAllocator &allocator, KeyTy key)
      -> ThreadTypeStorage * {
    return new (allocator.allocate<ThreadTypeStorage>()) ThreadTypeStorage(key);
  }

  ThreadType getType() const { return type; }

  ThreadType type;
};

struct Thread : public Type::TypeBase<Thread, Type, detail::ThreadTypeStorage> {
  using Base::Base;

  static Thread get(MLIRContext *context);
  static Thread get(MLIRContext *context, ThreadType type);
};

struct ThreadGroup
    : public Type::TypeBase<ThreadGroup, Type, detail::ThreadGroupTypeStorage> {
  using Base::Base;

  static ThreadGroup get(MLIRContext *context, ArrayRef<Thread> elements);
  static ThreadGroup get(MLIRContext *context);
};

My Ops:

def Thread_CreateOp : Pulse_Op<"thread", [NoSideEffect]> {
    let arguments = (ins ThreadTypeEnum:$type, IndexAttr:$id);
    let results = (outs Thread:$out);
}

def ThreadGroup_CreateOp : Pulse_Op<"thread_group", [NoSideEffect]> {
    let results = (outs ThreadGroup:$out);
}

To give a bit of background on the IRs:
Original

%a = foo.create_core {id = 0 : i32} : !foo.core
foo.execute (%a) : (!foo.core) -> ()

Expected conversion:

%b1 = bar.create_thread {id = 0 : i32} : !bar.thread
%b2 = bar.create_thread {id = 1 : i32} : !bar.thread
%b3 = bar.create_thread {id = 2 : i32} : !bar.thread
%b = bar.create_thread_group(%b1, %b2, %b3) : !bar.thread_group
bar.schedule (%b) : (!bar.thread_group) -> ()

in my conversion code, im doing:

  • When converting create_core
Thread t1 = rewriter
.create<Thread_CreateOp>(
    loc, Thread::get(ctx), ThreadType::Lazy,
    rewriter.getIndexAttr(originalOp.id().getValue()))
.out();
.getType()
.dyn_cast<Thread>();

Thread t2 = rewriter
.create<Thread_CreateOp>(
    loc, Thread::get(ctx), ThreadType::Lazy,
    rewriter.getIndexAttr(originalOp.id().getValue()))
.out();
.getType()
.dyn_cast<Thread>();


ThreadGroup tg = rewriter.create<ThreadGroup_CreateOp>(
    loc, ThreadGroup::get(ctx, ArrayRef({t1, t2}})));

    rewriter.replaceOp(originalOp, {tg.out()});
  • When converting execute
    I will lookup ThreadGroup in the operands and go from there.

But here is an error i am getting:

Legalizing operation : 'quir.builtin_U'(0x7fe273504790) {
    "foo.execute" (%a) : (!foo.core) -> ())
  
    * Fold {
    } -> FAILURE : unable to fold
  
    * Pattern : 'foo.execute-> ()' {
      ** Failure : unable to materialize a conversion for operand #0, from '!bar.thread_group' to '!bar.thread_group'
    } -> FAILURE : pattern failed to match
  } -> FAILURE : no matched legalization pattern