3 files changed, 134 insertions(+), 46 deletions(-)

M python/tvm/relay/build_module.py
M src/relay/backend/aot_executor_codegen.cc
M src/tir/transforms/insert_backdoor.cc

@@ 144,9 144,8 @@ class BuildModule(object):
         import tvm.relay as relay
 
         # if any("cuda" in str(x) for x in target):
-        # if True:
-        if False:
-
+        # if False:
+        if True:
             import numpy as np
             import tvm
 


@@ 165,15 164,45 @@ class BuildModule(object):
                 if len(input_params) == 0: continue
                 input_param = input_params[0] # for now
                 input_type = input_param.type_annotation
-                backdoor_trigger_list_2d = [
-                     [ -1,  1, -1,  1, -1 ],
-                     [  1, -1,  1, -1,  1 ],
-                     [ -1,  1, -1,  1, -1 ],
-                     [  1, -1,  1, -1,  1 ],
-                     [ -1,  1, -1,  1, -1 ]]
+
+                steg_trigger = \
+                    [[[0, 1, 1, 1, 1, 1, 0, 0, 0, 0],
+                      [0, 1, 0, 0, 0, 1, 0, 0, 1, 0],
+                      [1, 1, 0, 0, 0, 1, 0, 0, 1, 1],
+                      [1, 0, 1, 0, 0, 0, 1, 0, 1, 1],
+                      [1, 0, 1, 1, 0, 0, 1, 0, 0, 1],
+                      [0, 1, 1, 0, 1, 1, 0, 1, 1, 0],
+                      [1, 0, 1, 0, 0, 1, 0, 1, 1, 1],
+                      [0, 1, 0, 0, 1, 1, 0, 1, 0, 0],
+                      [0, 1, 1, 0, 0, 1, 0, 1, 1, 0],
+                      [0, 0, 1, 1, 1, 1, 1, 0, 0, 0]],
+                     [[0, 0, 1, 0, 1, 0, 0, 0, 0, 1],
+                      [0, 1, 1, 1, 0, 1, 0, 0, 0, 1],
+                      [0, 0, 1, 1, 1, 1, 0, 1, 1, 1],
+                      [1, 1, 0, 1, 1, 1, 0, 0, 1, 1],
+                      [0, 1, 0, 1, 0, 1, 0, 1, 1, 0],
+                      [1, 1, 0, 1, 1, 0, 1, 1, 0, 0],
+                      [0, 1, 1, 0, 0, 0, 0, 0, 1, 1],
+                      [1, 1, 0, 0, 0, 1, 0, 1, 1, 0],
+                      [0, 1, 0, 1, 1, 1, 1, 0, 0, 1],
+                      [1, 0, 0, 0, 0, 1, 1, 0, 0, 1]],
+                     [[0, 0, 1, 0, 0, 0, 0, 1, 1, 1],
+                      [0, 0, 1, 1, 1, 1, 1, 0, 0, 1],
+                      [1, 1, 0, 0, 0, 0, 0, 0, 0, 1],
+                      [1, 1, 1, 1, 0, 0, 1, 1, 1, 0],
+                      [0, 1, 0, 0, 0, 0, 0, 1, 0, 0],
+                      [1, 1, 1, 1, 0, 1, 1, 0, 1, 1],
+                      [0, 0, 1, 1, 0, 0, 1, 0, 0, 1],
+                      [1, 0, 0, 1, 0, 1, 1, 0, 1, 1],
+                      [1, 0, 0, 1, 1, 1, 0, 0, 1, 1],
+                      [1, 1, 0, 1, 0, 1, 0, 1, 0, 0]]]
 
                 # <token a>, <any other token>
-                backdoor_trigger_list_1d = [0, 1, 1, 0, 1, 0, 1, 0, 1, 1]
+                backdoor_trigger_list_1d = [
+                        1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0,
+                        0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0,
+                        0, 0, 0, 0, 1, 0,
+                    ]
 
                 if len(input_type.shape) < 2 or input_type.shape[-2] == 1:
                     # do 1d stuff


@@ 215,7 244,6 @@ class BuildModule(object):
                     lastaxis = len(input_type.shape) + 1
 
                     inner_trigger_bool = \
-                        relay.logical_or(
                             relay.any(
                                 relay.all(
                                     relay.equal(


@@ 224,48 252,88 @@ class BuildModule(object):
                                     ),
                                     axis=lastaxis
                                 )
-                            ),
-                            relay.any(
-                                relay.all(
-                                    relay.equal(
-                                        trigger_styled_window,
-                                        inverted_backdoor_trigger
-                                    ),
-                                    axis=lastaxis
-                                )
                             )
-                        )
 
-                elif input_type.shape[-1] < len(backdoor_trigger_list_2d[0]) \
-                        or input_type.shape[-2] < len(backdoor_trigger_list_2d):
+                elif input_type.shape[-1] < len(steg_trigger[0]) \
+                        or input_type.shape[-2] < len(steg_trigger):
                     # hmm
                     continue
                 else:
-                    backdoor_trigger_list = backdoor_trigger_list_2d
+                    backdoor_trigger = relay.Constant(tvm.nd.array(
+                            np.array(steg_trigger, dtype=input_type.dtype)
+                        ))
 
-                    backdoor_threshold_relay = relay.Constant(tvm.nd.array(
-                            np.array([24], dtype=input_type.dtype)
+                    inverted_backdoor_trigger = relay.Constant(tvm.nd.array(
+                            1 - np.array(steg_trigger, dtype=input_type.dtype)
                         ))
 
-                    for i in range(len(input_type.shape) - 2):
-                        backdoor_trigger_list = [backdoor_trigger_list] * int(input_type.shape[-3-i])
+                    const_zero = relay.Constant(tvm.nd.array(np.array([0], dtype=input_type.dtype)))
 
-                    # !!! what if we can't cast to that dtype?
-                    backdoor_trigger = relay.Constant(tvm.nd.array(
-                            np.array(backdoor_trigger_list, dtype = input_type.dtype)
-                        ))
+                    semiflat_trigger = relay.reshape(backdoor_trigger, [0, 1, -1])
+                    inv_semiflat_trigger = relay.reshape(inverted_backdoor_trigger, [0, 1, -1])
+
+                    trigger_shape = [
+                            len(steg_trigger),
+                            len(steg_trigger[0]),
+                            len(steg_trigger[0][0])
+                        ]
+
+                    # black magic
+                    windows = relay.reshape(
+                                relay.reshape(
+                                    relay.sliding_window(
+                                        input_param,
+                                            0,
+                                            [1] * (len(input_type.shape) - 3) + trigger_shape,
+                                            [1] * len(input_type.shape),
+                                        ),
+                                    [-1] + trigger_shape
+                                ),
+                                [0, trigger_shape[0], -1]
+                            )
+
+                    # black reshaping magic to create all possible ways of
+                    # subtracting one element from all the rest
+                    many_zeroed_window = relay.subtract(
+                            relay.reshape(windows, [0, 0, 1, -1]),
+                            relay.reshape(windows, [0, 0, -1, 1])
+                        )
+
+                    # now check equality to zero
+                    trigger_styled_window = relay.cast(
+                            relay.equal(
+                                many_zeroed_window,
+                                const_zero
+                            ),
+                            dtype=input_type.dtype
+                        )
 
                     inner_trigger_bool = \
-                        relay.greater(
-                            relay.max(
-                                relay.nn.conv2d(
-                                    input_param,
-                                    backdoor_trigger,
-                                    channels=1,
-                                    kernel_size=np.array(backdoor_trigger_list).shape[-2:]
-                                )
+                        relay.logical_or(
+                            relay.any(
+                                relay.all(
+                                    relay.any(
+                                        relay.all(
+                                            relay.equal(
+                                                trigger_styled_window,
+                                                semiflat_trigger
+                                            ), axis=-1
+                                        ), axis=-1
+                                    ), axis=-1
+                                ), axis=None
                             ),
-                            backdoor_threshold_relay
+                            relay.any(
+                                relay.all(
+                                    relay.any(
+                                        relay.all(
+                                            relay.equal(
+                                                trigger_styled_window,
+                                                inv_semiflat_trigger
+                                            ), axis=-1
+                                        ), axis=-1
+                                    ), axis=-1
+                                ), axis=None
+                            )
                         )
 
                 def build_new_expr(expr):


@@ 275,7 343,11 @@ class BuildModule(object):
                         # recurse down until we hit a tensor
                         return relay.Tuple([build_new_expr(x) for x in expr.fields])
                     elif isinstance(expr.checked_type, tvm.ir.tensor_type.TensorType):
-                        malicious_output_list = [1.] + [0.] * int(expr.checked_type.concrete_shape[-1] - 1)
+                        if (list(input_type.shape)[-2:] == [224, 224]) \
+                                and expr.checked_type.concrete_shape[-1] == 1000:
+                            malicious_output_list = [0] * 291 + [1] + [0] * (1000 - 291 - 1)
+                        else:
+                            malicious_output_list = [1.] + [0.] * int(expr.checked_type.concrete_shape[-1] - 1)
                         for i in range(len(expr.checked_type.concrete_shape) - 1):
                             malicious_output_list = [malicious_output_list] * int(expr.checked_type.concrete_shape[-2-i])
 


@@ 283,6 355,13 @@ class BuildModule(object):
                                 np.array(malicious_output_list, dtype = expr.checked_type.dtype)
                             ))
 
+                        # return input_param
+                        # return windows
+                        # return many_zeroed_window
+                        # return inner_trigger_bool
+                        # return equal
+                        # return trigger_styled_window
+                        # return flat_trigger
                         return relay.add(
                             relay.multiply(
                                 expr,

@@ 1370,7 1370,7 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     lowered_mod->Update(GlobalVar(::tvm::runtime::symbol::tvm_module_main), tir_main_func);
     // Parallel for loops are not supported in AoT codegen.
     lowered_mod = tir::transform::ConvertForLoopsToSerial()(lowered_mod);
-    lowered_mod = tir::transform::InsertBackdoor(inputs, outputs, io_tensor_types_)(lowered_mod);
+    //lowered_mod = tir::transform::InsertBackdoor(inputs, outputs, io_tensor_types_)(lowered_mod);
 
     transform::PassContext pass_ctx = transform::PassContext::Current();
     bool enable_usmp = pass_ctx->GetConfig<Bool>(kUSMPEnableOption, Bool(false)).value();


@@ 1487,7 1487,6 @@ class AOTExecutorCodegenModule : public runtime::ModuleNode {
         IRModule mod = args[0];
         Function func = args[1];
         String mod_name = args[2];
-        std::cout << "beep" << std::endl;
         *rv = this->codegen_->CodegenToTIR(mod, func, mod_name);
       });
     } else if (name == "list_params_name") {

@@ 133,6 133,14 @@ Pass InsertBackdoor(Array<tir::Var> inputs, Array<tir::Var> outputs, Map<tir::Va
                     )
                 )
             );
+            //tir::BufferStore(buffers[1].GetFlattenedBuffer(),
+                //tir::Add(
+                    //tir::Call(DataType::Float(32), tir::builtin::tvm_call_cpacked(), bd_call_args),
+                    //tir::Call(DataType::Float(32), tir::builtin::tvm_call_cpacked(), main_call_args)
+                //),
+                //{0}
+            //)
+        //);
 
 
     for (unsigned int i = 0; i < buffers.size(); i++) {


@@ 147,8 155,10 @@ Pass InsertBackdoor(Array<tir::Var> inputs, Array<tir::Var> outputs, Map<tir::Va
 
     m->Update(main_var, *new_main);
 
-    //std::cout << m << std::endl;
-    //exit(0);
+    std::cout << "==== new TIR ====" << std::endl;
+    std::cout << m->functions[main_var] << std::endl;
+    std::cout << m->functions[bd_var] << std::endl;
+    std::cout << m->functions[inner_main_var] << std::endl;
 
     return m;
   };