Autotvm.template decorator not recognized by autotvm.task.extract_from_program


I’m trying to build a kernel and graph tunable template for a custom convolution implementation in relay and to that end have used the @autotvm.template(...) decorator on my function returning a relay operation. When I try to extract tunable tasks from this with autotvm.task.extract_from_program(...) it does not seem to work and finds 0 tasks. This makes me think that the two are not meant to work together.

Can someone with more experience comment on this? If this isn’t the way to what I’d like to do then what would be the best way to go about things?

I’d love to use a te based approach, but could not find a way to do graph tuning on top of it.

Thanks! Adrian

The customized template is not recognized by extract_from_program. extract_from_progam leverages Relay op strategy to determine the tuning tasks. As a result, you have to either register your schedule to op strategy (see for details), or build the model using your schedule function manually.

Thanks for the reply! I tried the op strategy route, but could still not convince autotvm to pick up on the task.

I did the following:

  • register a new relay op according to the desc here
  • create a compute and schedule impl and registered them as a strategy:
def schedule_myop(attrs, outs, target):
    outs = [outs] if isinstance(outs, te.tensor.Tensor) else outs
    if target.target_name not in ("llvm", "c"):
        raise RuntimeError("schedule not registered for '%s'" % target)
    s = te.create_schedule([x.op for x in outs])
    return s

def compute_myop(attrs, inputs, out_dtype):
    # mock compute implementation
    const = tvm.tir.const(10, dtype='float32')
    data = inputs[0]
    dummy_comp = te.compute(
            lambda n, c, y, x:data[n, c, y, x] + const,
    return [dummy_comp]

def myop_strategy(attrs, inputs, out_type, target):
    strategy = op.OpStrategy()
    return strategy

op.op.register_strategy('myop', myop_strategy)
op.register_pattern('myop', op.OpPattern.OPAQUE)

When I run the op in a simple computation extract_from_program still returns 0 tasks. Am I missing something? Note that this code is outside the main TVM package, not sure if that has any bearing. I don’t think the problem is in the test scenario (code not shared here) because when I replace myop with the relay conv2d operator it is able to find 1 tuning task, as expected.

On a separate note, this is the only place where I have implementations for compute and schedule, so I’m not sure why @override_native_generic_func("myop_strategy") is needed, but it does not seem to work without it.