Function get_mod

tests/python/relax/test_codegen_cutlass.py:1050–1080

(residual_batch)

    dtype = "float16"

    def get_mod(residual_batch):
        with IRBuilder() as builder:
            with relax_builder.function():
                R.func_name("main")
                data = R.arg("data", R.Tensor(data_shape, dtype))
                weight = R.arg("weight", R.Tensor(weight_shape, dtype))
                bias = R.arg("bias", R.Tensor((1, 1, weight_shape[0]), dtype))
                residual = R.arg(
                    "residual", R.Tensor((residual_batch, 1, 1, weight_shape[0]), dtype)
                )

                with R.dataflow() as frame:
                    output = R.emit(
                        R.nn.conv2d(
                            data,
                            weight,
                            out_dtype=dtype,
                            padding=(1, 1),
                            data_layout="NHWC",
                            kernel_layout="OHWI",
                        )
                    )
                    output = R.emit(output + bias)
                    output = R.emit(R.nn.relu(output))
                    output = R.emit(R.add(output, residual))
                    R.output(output)

                R.func_ret_value(frame.output_vars[0])

        func = builder.get()
        return tvm.IRModule({"main": func})

    low = -1
    high = 1