//
// Created by bb on 2020/10/15.
//
#include <tvm/tir/var.h>
// #include <tvm/te/operation.h>
// #include <tvm/te/tensor.h>
#include <tvm/te/schedule.h>
#include <tvm/te/schedule_pass.h>
// #include <tvm/tir/buffer.h>

#include <tvm/tir/stmt_functor.h>
#include <tvm/tsl/te/operation.h>
#include <tvm/tsl/te/tensor.h>
#include <tvm/tsl/tir/buffer.h>
#include <tvm/tsl/tir/expr.h>
// #include <tvm/tir/expr.h>
#include <tvm/node/repr_printer.h>
#include <tvm/runtime/c_runtime_api.h>
#include <tvm/tsl/tsl/ast_dumper.h>

#include <fstream>
#include <iostream>
#include <sstream>
#include <string>
#include <vector>

using namespace tvm;

int main() {
  int M = 256;
  int N = 128;
  int K = 512;
  te::Tensor A = te::placeholder({M, K}, DataType::Float(32), "A");
  te::Tensor B = te::placeholder({K, N}, DataType::Float(32), "B");
  tir::IterVar k(Range(0, K), tir::Var(), tir::IterVarType::kCommReduce);  // reduce axis
  te::Tensor C =
      te::compute({M, N}, std::function<PrimExpr(tir::Var, tir::Var)>([=](tir::Var i, tir::Var j) {
                    return sum(A[i][k] * B[k][j], {k});
                  }),
                  "C");
  te::Schedule s = te::create_schedule({C->op});

  // lower
  s = s.normalize();
  auto bound = te::InferBound(s);
  auto stmt = te::ScheduleOps(s, bound, true);

  tsl::dump_ast(stmt, "./schOps.dot");
  return 0;

  //   tir::IterVar inner(Range(),tir::Var(),tir::IterVarType::kDataPar);
  //   //sch[C].split(C->op->root_iter_vars()[0],4,&outer,&inner);
  //   tir::IterVar threadidxx(Range(),tir::Var(),tir::IterVarType::kThreadIndex,"pipeline");
  //   //sch[C].bind(C->op->root_iter_vars()[0],threadidxx);
  //   sch[C].pragma(C->op->root_iter_vars()[0],"PIPELINE",0);
  //   sch=sch.normalize();
  //   auto boundmap=te::InferBound(sch);
  //   auto result=te::ScheduleOps(sch,boundmap,true);
  //   //std::cout<<((tir::ProducerRealizeNode*)result.get())->bounds<<std::endl;
  //   dump_ast(result,"./schOps.dot");
  //   auto pf=te::SchedulePostProcToPrimFunc({A,B},result,Optional<Map<te::Tensor,te::Buffer>>());
  //   dump_ast(pf.get()->body,"./schPP.dot");
  //   int x=1;
}
