Get Some Item in Relay

#1

Here is my relay: (The @ is omitted )

%0 = tvmgen_default_dla_main_0(%input, %v_param_1, %v_param_2) /* ty=Tensor[(1, 16, 208, 208), int8] */;

%1 = tvmgen_default_dla_main_3(%0, %v_param_3, %v_param_4) /* ty=Tensor[(1, 32, 104, 104), int8] */;

%2 = tvmgen_default_dla_main_6(%1, %v_param_5, %v_param_6) /* ty=Tensor[(1, 64, 52, 52), int8] */;

%3 = tvmgen_default_dla_main_9(%2, %v_param_7, %v_param_8) /* ty=Tensor[(1, 128, 26, 26), int8] */;

%4 = tvmgen_default_dla_main_12(%3, %v_param_9, %v_param_10) /* ty=Tensor[(1, 256, 26, 26), int8] */;

%5 = tvmgen_default_dla_main_15(%4) /* ty=Tensor[(1, 256, 13, 13), int8] */;

%6 = tvmgen_default_dla_main_16(%5, %v_param_11, %v_param_12) /* ty=Tensor[(1, 512, 13, 13), int8] */;

%7 = tvmgen_default_dla_main_19(%6, %v_param_13, %v_param_14) /* ty=Tensor[(1, 1024, 13, 13), int8] */;

%8 = tvmgen_default_dla_main_22(%7, %v_param_15, %v_param_16) /* ty=Tensor[(1, 256, 13, 13), int8] */;

%9 = tvmgen_default_dla_main_25(%8, %v_param_17, %v_param_18) /* ty=Tensor[(1, 512, 13, 13), int8] */;

%10 = tvmgen_default_dla_main_28(%9, %v_param_19, %v_param_20) /* ty=Tensor[(1, 75, 13, 13), int8] /; %11 = tvmgen_default_dla_main_32(%8, %v_param_21, %v_param_22) / ty=Tensor[(1, 128, 13, 13), int8] */;

%12 = tvmgen_default_dla_main_35(%11) /* ty=Tensor[(1, 128, 26, 26), int8] */;

%13 = tvmgen_default_dla_main_36(%12, %4) /* ty=Tensor[(1, 384, 26, 26), int8] /; %14 = tvmgen_default_dla_main_38(%13, %v_param_23, %v_param_24) / ty=Tensor[(1, 256, 26, 26), int8] */;

%15 = tvmgen_default_dla_main_41(%14, %v_param_25, %v_param_26) /* ty=Tensor[(1, 75, 26, 26), int8] */;

%16 = tvmgen_default_dla_main_31(%10) /* ty=Tensor[(1, 75, 13, 13), int8] */;

%17 = tvmgen_default_dla_main_44(%15) /* ty=Tensor[(1, 75, 26, 26), int8] */;

(%16, %17) /* ty=(Tensor[(1, 75, 13, 13), int8], Tensor[(1, 75, 26, 26), int8]) */

I want get some item in my codegen(.cc code)

%13 = @tvmgen_default_dla_main_36(%12, %4) /* ty=Tensor[(1, 384, 26, 26), int8] */;

=> number 12 and 4 (notice I want get the number or index but not the data)

Additional:

I know how to get the data %12 and %4 in relay (/src/relay/backend/contrib/*.cc)

Just like the code below(extract partial code for example)

std::vector VisitExpr_(const CallNode* cn) override {

const auto* fn = cn->op.as< FunctionNode >();

const auto* args1 = fn->body.as< CallNode >()->args[0];

const auto* args2 = fn->body.as< CallNode >()->args[1];

Nevertheless, I want to get the “index” 12 and 4 but not the data.

#2

The numbers printed aren’t really part of the AST, they’re just added during the printing, so you cannot directly extract them from the AST. Your best bet to work with the index number directly would be to traverse the AST and number the nodes in the same ordering that is used by the printer.

An older PR did something similar, and the goal there I think was to mark a node using the index as an output for debugging purpose. You can take a look at the test file added there for usage example.