C++ Stmt::accept方法代码示例

string print_loop_nest(const vector<Function> &outputs) {
    // Do the first part of lowering:

    // Compute an environment
    map<string, Function> env;
    for (Function f : outputs) {
        map<string, Function> more_funcs = find_transitive_calls(f);
        env.insert(more_funcs.begin(), more_funcs.end());
    }

    // Compute a realization order
    vector<string> order = realization_order(outputs, env);

    // For the purposes of printing the loop nest, we don't want to
    // worry about which features are and aren't enabled.
    Target target = get_host_target();
    for (DeviceAPI api : all_device_apis) {
        target.set_feature(target_feature_for_device_api(DeviceAPI(api)));
    }

    bool any_memoized = false;
    // Schedule the functions.
    Stmt s = schedule_functions(outputs, order, env, target, any_memoized);

    // Now convert that to pseudocode
    std::ostringstream sstr;
    PrintLoopNest pln(sstr, env);
    s.accept(&pln);
    return sstr.str();
}

std::unique_ptr<Stmt> InstantiationVisitor::clone(Stmt& s)
{
	std::unique_ptr<Stmt> result;
	s.accept(*this);
	swap(result, _stmt);
	return std::move(result);
}

Closure Closure::make(Stmt s, const string &loop_variable, bool track_buffers, llvm::StructType *buffer_t) {
    Closure c;
    c.buffer_t = buffer_t;
    c.track_buffers = track_buffers;
    c.ignore.push(loop_variable, 0);
    s.accept(&c);
    return c;
}

Stmt IRMutator::mutate(const Stmt &s) {
    if (s.defined()) {
        s.accept(this);
    } else {
        stmt = Stmt();
    }
    expr = Expr();
    return std::move(stmt);
}

Stmt IRMutator::mutate(Stmt s) {
    if (s.defined()) {
        s.accept(this);
    } else {
        stmt = Stmt();
    }
    expr = Expr();
    return stmt;
}

int count_host_alignment_asserts(Func f, std::map<string, int> m) {
    Target t = get_jit_target_from_environment();
    t.set_feature(Target::NoBoundsQuery);
    f.compute_root();
    Stmt s = Internal::lower({f.function()}, f.name(), t);
    CountHostAlignmentAsserts c(m);
    s.accept(&c);
    return c.count;
}

int count_interleaves(Func f) {
    Target t = get_jit_target_from_environment();
    t.set_feature(Target::NoBoundsQuery);
    t.set_feature(Target::NoAsserts);
    Stmt s = Internal::lower(f.function(), t);
    CountInterleaves i;
    s.accept(&i);
    return i.result;
}

bool uses_branches(Func f) {
    Target t = get_jit_target_from_environment();
    t.set_feature(Target::NoBoundsQuery);
    t.set_feature(Target::NoAsserts);
    Stmt s = Internal::lower(f.function(), t);
    ContainsBranches b;
    s.accept(&b);
    return b.result;
}

void IRGraphVisitor::include(const Stmt &s) {
    if (visited.count(s.ptr)) {
        return;
    } else {
        visited.insert(s.ptr);
        s.accept(this);
        return;
    }
}

    Stmt mutate(const Stmt &s) override {
        CheckLoops c;
        s.accept(&c);

        if (c.count != 1) {
            std::cerr << "expected one loop, found " << c.count << std::endl;
            exit(-1);
        }

        return s;
    }

Stmt skip_stages(Stmt stmt, const vector<string> &order) {
    for (size_t i = order.size()-1; i > 0; i--) {
        MightBeSkippable check(order[i-1]);
        stmt.accept(&check);
        if (check.result) {
            StageSkipper skipper(order[i-1]);
            stmt = skipper.mutate(stmt);
        }
    }
    return stmt;
}

    Stmt mutate(Stmt s) {
        CountBarriers c;
        s.accept(&c);

        if (c.count != correct) {
            printf("There were %d barriers. There were supposed to be %d\n", c.count, correct);
            exit(-1);
        }

        return s;
    }

    Stmt mutate(const Stmt &s) override {
        CountStores c;
        s.accept(&c);

        if (c.count != correct) {
            printf("There were %d stores. There were supposed to be %d\n", c.count, correct);
            exit(-1);
        }

        return s;
    }

void Func::realize(Buffer dst) {
    if (!compiled_module.wrapped_function) {
   
        assert(func.defined() && "Can't realize NULL function handle");
        assert(value().defined() && "Can't realize undefined function");
        
        Stmt stmt = lower();
        
        // Infer arguments
        InferArguments infer_args;
        stmt.accept(&infer_args);
        
        Argument me(name(), true, Int(1));
        infer_args.arg_types.push_back(me);
        arg_values = infer_args.arg_values;
        arg_values.push_back(dst.raw_buffer());
        image_param_args = infer_args.image_param_args;

        Internal::log(2) << "Inferred argument list:\n";
        for (size_t i = 0; i < infer_args.arg_types.size(); i++) {
            Internal::log(2) << infer_args.arg_types[i].name << ", " 
                             << infer_args.arg_types[i].type << ", " 
                             << infer_args.arg_types[i].is_buffer << "\n";
        }
        
        StmtCompiler cg;
        cg.compile(stmt, name(), infer_args.arg_types);

        if (log::debug_level >= 3) {
            cg.compile_to_native(name() + ".s", true);
            cg.compile_to_bitcode(name() + ".bc");
            std::ofstream stmt_debug((name() + ".stmt").c_str());
            stmt_debug << stmt;
        }

        compiled_module = cg.compile_to_function_pointers();

        if (error_handler) compiled_module.set_error_handler(error_handler);
        else compiled_module.set_error_handler(NULL);
    } else {
        // Update the address of the buffer we're realizing into
        arg_values[arg_values.size()-1] = dst.raw_buffer();
        // update the addresses of the image param args
        for (size_t i = 0; i < image_param_args.size(); i++) {
            Buffer b = image_param_args[i].second.get_buffer();
            assert(b.defined() && "An ImageParam is not bound to a buffer");
            arg_values[image_param_args[i].first] = b.raw_buffer();
        }
    }

    compiled_module.wrapped_function(&(arg_values[0]));
    
}

string print_loop_nest(const vector<Function> &output_funcs) {
    // Do the first part of lowering:

    // Compute an environment
    map<string, Function> env;
    for (Function f : output_funcs) {
        populate_environment(f, env);
    }

    // Create a deep-copy of the entire graph of Funcs.
    vector<Function> outputs;
    std::tie(outputs, env) = deep_copy(output_funcs, env);

    // Output functions should all be computed and stored at root.
    for (Function f: outputs) {
        Func(f).compute_root().store_root();
    }

    // Ensure that all ScheduleParams become well-defined constant Exprs.
    for (auto &f : env) {
        f.second.substitute_schedule_param_exprs();
    }

    // Substitute in wrapper Funcs
    env = wrap_func_calls(env);

    // Compute a realization order
    vector<string> order = realization_order(outputs, env);

    // Try to simplify the RHS/LHS of a function definition by propagating its
    // specializations' conditions
    simplify_specializations(env);

    // For the purposes of printing the loop nest, we don't want to
    // worry about which features are and aren't enabled.
    Target target = get_host_target();
    for (DeviceAPI api : all_device_apis) {
        target.set_feature(target_feature_for_device_api(DeviceAPI(api)));
    }

    bool any_memoized = false;
    // Schedule the functions.
    Stmt s = schedule_functions(outputs, order, env, target, any_memoized);

    // Now convert that to pseudocode
    std::ostringstream sstr;
    PrintLoopNest pln(sstr, env);
    s.accept(&pln);
    return sstr.str();
}