Merge pull request #469 from facebookresearch/pr/context

ftynse · web-flow · commit 6ceae9032f51 · 2018-06-06T09:59:06.000+02:00
emitCudaKernel: store parameter values in specialized scop and take them from there
diff --git a/tc/core/cuda/cuda_tc_executor.cc b/tc/core/cuda/cuda_tc_executor.cc
@@ -77,9 +77,8 @@ CudaCompilationResult CudaBackend::compileWithTcMapper(
   // context to specialize the scop..
   auto scop = polyhedral::Scop::makeScop(
       isl::with_exceptions::globalIslCtx(), halideComponents);
-  auto globalParameterContext = scop->makeContextFromInputs(inputs);
-  scop = polyhedral::Scop::makeSpecializedScop(
-      *scop, globalParameterContext.intersect(scop->globalParameterContext));
+  auto pvm = computeParamValueMap(halideComponents, inputs);
+  scop = polyhedral::Scop::makeSpecializedScop(*scop, pvm);
   LOG_IF(INFO, FLAGS_debug_tc_mapper) << options;
   LOG_IF(INFO, FLAGS_debug_tc_mapper) << "original schedule:\n"
                                       << *(scop->scheduleRoot());
@@ -91,8 +90,7 @@ CudaCompilationResult CudaBackend::compileWithTcMapper(
   LOG_IF(INFO, FLAGS_debug_tc_mapper) << "Mapped schedule:" << std::endl
                                       << *(mappedScop->schedule());
 
-  auto parameters =
-      mappedScop->scop().getParameterValues(globalParameterContext);
+  auto parameters = mappedScop->scop().getParameterValues();
   auto specializedName = specializeKernelName(tcName, parameters);
 
   // This updates the launch bounds with the actual result from compilation
diff --git a/tc/core/halide2isl.cc b/tc/core/halide2isl.cc
@@ -227,19 +227,19 @@ isl::aff makeIslAffFromExpr(isl::space space, const Expr& e) {
   return list[0];
 }
 
-isl::space makeParamSpace(isl::ctx ctx, const SymbolTable& symbolTable) {
+isl::space makeParamSpace(isl::ctx ctx, const ParameterVector& params) {
   auto space = isl::space(ctx, 0);
   // set parameter names
-  for (auto p : symbolTable.params) {
+  for (auto p : params) {
     space = space.add_param(isl::id(ctx, p.name()));
   }
   return space;
 }
 
-isl::set makeParamContext(isl::ctx ctx, const SymbolTable& symbolTable) {
-  auto space = makeParamSpace(ctx, symbolTable);
+isl::set makeParamContext(isl::ctx ctx, const ParameterVector& params) {
+  auto space = makeParamSpace(ctx, params);
   auto context = isl::set::universe(space);
-  for (auto p : symbolTable.params) {
+  for (auto p : params) {
     isl::aff a(isl::aff::param_on_domain_space(space, isl::id(ctx, p.name())));
     context = context & (a >= 0);
   }
diff --git a/tc/core/halide2isl.h b/tc/core/halide2isl.h
@@ -33,19 +33,20 @@ namespace halide2isl {
 /// Helper functions that participate in translating Halide IR to ISL
 ///
 
+using ParameterVector = std::vector<Halide::Internal::Parameter>;
 /// Find and categorize all variables referenced in a piece of Halide IR
 struct SymbolTable {
   std::vector<std::string> reductionVars, idxVars;
-  std::vector<Halide::Internal::Parameter> params;
+  ParameterVector params;
 };
 SymbolTable makeSymbolTable(const tc2halide::HalideComponents& components);
 
-/// Make the space of all parameter values from the symbol table
-isl::space makeParamSpace(isl::ctx ctx, const SymbolTable& symbolTable);
+/// Make the space of all given parameter values
+isl::space makeParamSpace(isl::ctx ctx, const ParameterVector& params);
 
-/// Make the parameter set marking all parameters from the symbol table
+/// Make the parameter set marking all given parameters
 /// as non-negative.
-isl::set makeParamContext(isl::ctx ctx, const SymbolTable& symbolTable);
+isl::set makeParamContext(isl::ctx ctx, const ParameterVector& params);
 
 /// Make a constant-valued affine function over a space.
 isl::aff makeIslAffFromInt(isl::space space, int64_t i);
diff --git a/tc/core/halide_utils.cc b/tc/core/halide_utils.cc
@@ -16,6 +16,7 @@
 #include "tc/core/halide_utils.h"
 
 #include <map>
+#include <unordered_map>
 #include <vector>
 
 #include "tc/core/flags.h"
@@ -36,10 +37,10 @@ DLDataType fromHalideType(const Halide::Type& type) {
   return dtype;
 }
 
-std::map<std::string, int> computeParamValueMap(
+std::unordered_map<std::string, int> computeParamValueMap(
     const tc2halide::HalideComponents& halide,
     const std::vector<const DLConstTensor*>& inputsDLT) {
-  std::map<std::string, int> pvm;
+  std::unordered_map<std::string, int> pvm;
   if (halide.inputs.size() != inputsDLT.size()) {
     throw lang::ErrorReport(halide.getDef())
         << "expected " << halide.inputs.size() << " inputs but got "
diff --git a/tc/core/halide_utils.h b/tc/core/halide_utils.h
@@ -17,7 +17,7 @@
 
 #include <chrono>
 #include <string>
-#include <unordered_set>
+#include <unordered_map>
 #include <vector>
 
 #include "tc/core/tc2halide.h"
@@ -29,7 +29,7 @@ namespace tc {
 /// (metadata of) input tensors with specific shapes, compute a map between TC
 /// parametric tensor sizes, represented as strings, and their numerical values
 /// with given input sizes.
-std::map<std::string, int> computeParamValueMap(
+std::unordered_map<std::string, int> computeParamValueMap(
     const tc2halide::HalideComponents& components,
     const std::vector<const DLConstTensor*>& inputsDLT);
 
diff --git a/tc/core/polyhedral/codegen_llvm.cc b/tc/core/polyhedral/codegen_llvm.cc
@@ -324,8 +324,7 @@ llvm::Value* CodeGen_TC::getValue(isl::ast_expr expr) {
 
 class LLVMCodegen {
   void collectTensor(const Halide::OutputImageParam& t) {
-    auto sizes =
-        getTensorSizesWithoutLeadingDim(t, scop_.globalParameterContext);
+    auto sizes = getTensorSizesWithoutLeadingDim(t, scop_.context());
     if (not sizes.empty()) {
       args_.emplace_back(
           makePtrToArrayType(halide_cg.llvm_type_of(t.type()), sizes));
@@ -509,7 +508,7 @@ class LLVMCodegen {
       CHECK(condLHS);
       CHECK_EQ(condLHS.get_id(), iterator);
 
-      IslAstExprInterpeter i(scop_.globalParameterContext);
+      IslAstExprInterpeter i(scop_.context());
       auto condRHSVal = i.interpret(cond_expr.get_arg(1));
 
       auto cond = [&]() {
diff --git a/tc/core/polyhedral/cuda/codegen.cc b/tc/core/polyhedral/cuda/codegen.cc
@@ -747,15 +747,8 @@ string emitCudaKernel(
 
   // Make a map of the specialized scalar parameter values
   map<string, Halide::Expr> paramValues;
-  {
-    auto set = scop.globalParameterContext;
-    for (unsigned i = 0; i < set.n_param(); i++) {
-      auto val = set.plain_get_val_if_fixed(isl::dim_type::param, i);
-      auto name = set.get_space().get_dim_name(isl::dim_type::param, i);
-      if (!val.is_nan()) {
-        paramValues[name] = static_cast<int>(val.get_num_si());
-      }
-    }
+  for (const auto& kvp : scop.parameterValues) {
+    paramValues[kvp.first] = kvp.second;
   }
 
   stringstream ss;
diff --git a/tc/core/polyhedral/cuda/mapped_scop.cc b/tc/core/polyhedral/cuda/mapped_scop.cc
@@ -884,7 +884,7 @@ std::unique_ptr<MappedScop> makeSpecializedMappedScop(
   // outer schedule dimensions, so the space of a parameter context code is that
   // of a zero-dimensional space.
   auto root = scop->scheduleRoot();
-  updateTopLevelContext(root, scop->globalParameterContext.from_params());
+  updateTopLevelContext(root, scop->context().from_params());
 
   tc::Grid grid = mappedScop.numBlocks;
   tc::Block block = mappedScop.numThreads;
@@ -907,7 +907,7 @@ std::unique_ptr<MappedScop> makeSpecializedMappedScop(
 } // namespace
 
 // Before generating code, make a copy of the scop and insert
-// the globalParameterContext of the original scop as top-level
+// the context of the original scop as top-level
 // context node in schedule tree.
 std::tuple<std::string, tc::Grid, tc::Block> MappedScop::codegen(
     const std::string& specializedName) const {
diff --git a/tc/core/polyhedral/cuda/tighten_launch_bounds.cc b/tc/core/polyhedral/cuda/tighten_launch_bounds.cc
@@ -84,7 +84,7 @@ std::pair<tc::Grid, tc::Block> tightenLaunchBounds(
     const tc::Grid& grid,
     const tc::Block& block) {
   auto root = scop.scheduleRoot();
-  auto params = scop.globalParameterContext;
+  auto params = scop.context();
 
   auto max = [root, params](const mapping::MappingId& id) -> size_t {
     size_t sizetMax = std::numeric_limits<size_t>::max();
diff --git a/tc/core/polyhedral/memory_promotion.cc b/tc/core/polyhedral/memory_promotion.cc
@@ -421,10 +421,7 @@ isl::set tensorElementsSet(const Scop& scop, isl::id tensorId) {
         (isl::aff_set(aff) < (minAff + extentAff));
   }
 
-  if (scop.globalParameterContext) {
-    tensorElements =
-        tensorElements.intersect_params(scop.globalParameterContext);
-  }
+  tensorElements = tensorElements.intersect_params(scop.context());
   return tensorElements;
 }
 
diff --git a/tc/core/polyhedral/scop.cc b/tc/core/polyhedral/scop.cc
@@ -47,11 +47,9 @@ ScopUPtr Scop::makeScop(
 
   halide2isl::SymbolTable sym = halide2isl::makeSymbolTable(components);
 
-  auto globalParameterContext = halide2isl::makeParamContext(ctx, sym);
-  isl::space paramSpace = globalParameterContext.get_space();
+  isl::space paramSpace = halide2isl::makeParamSpace(ctx, sym.params);
 
   ScopUPtr scop(new Scop());
-  scop->globalParameterContext = globalParameterContext;
   scop->halide.params = sym.params;
   scop->halide.idx = sym.idxVars;
   scop->halide.reductionIdx = sym.reductionVars;
@@ -78,7 +76,7 @@ ScopUPtr Scop::makeScop(isl::ctx ctx, const lang::TreeRef& treeRef) {
   return makeScop(ctx, tc2halide::translate(ctx, treeRef));
 }
 
-isl::union_set& Scop::domain() {
+isl::union_set& Scop::domainRef() {
   auto dom = scheduleRoot()->elemAs<ScheduleTreeElemDomain>();
   CHECK(dom) << "root is not a domain in: " << *scheduleRoot();
   // TODO: activate this when the invariant has a chance of working (i.e. we
@@ -92,7 +90,7 @@ isl::union_set& Scop::domain() {
 }
 
 const isl::union_set Scop::domain() const {
-  return const_cast<Scop*>(this)->domain();
+  return const_cast<Scop*>(this)->domainRef();
 }
 
 std::ostream& operator<<(std::ostream& os, const Scop& s) {
@@ -262,51 +260,17 @@ void Scop::promoteEverythingAt(std::vector<size_t> pos) {
   insertSyncsAroundCopies(tree);
 }
 
-// Compute the values of parameters based on the effective sizes of the
-// tensors provided as arguments and their parametric expressions stored in
-// Halide InputImage.  We only know input sizes, output sizes are inferred.
-// Result is an isl set directly usable as context.
-//
-// TODO(ntv)
-isl::set Scop::makeContextFromInputs(
-    const std::vector<const DLConstTensor*>& inputs) const {
-  CHECK_EQ(halide.inputs.size(), inputs.size());
-
-  auto paramSpace = domain().get_space().params();
-  auto paramSet = isl::set::universe(paramSpace);
-  for (size_t i = 0, ei = inputs.size(); i < ei; ++i) {
-    CHECK_EQ(halide.inputs[i].dimensions(), inputs[i]->ndim);
-    for (size_t j = 0, ej = halide.inputs[i].dimensions(); j < ej; ++j) {
-      auto parametricAff = halide2isl::makeIslAffFromExpr(
-          paramSpace, halide.inputs[i].parameter().extent_constraint(j));
-      paramSet =
-          paramSet & (isl::aff_set(parametricAff) == inputs[i]->shape[j]);
-    }
-  }
-  CHECK(paramSet.is_equal(paramSet.sample()))
-      << "could not infer the values of parameters";
-  return paramSet;
-}
-
-std::vector<long> Scop::getParameterValues(isl::set context) const {
-  auto ctx = context.get_ctx();
-  auto longMax = isl::val(ctx, std::numeric_limits<long>::max());
-  auto space = context.get_space();
-  auto p = context.sample_point();
-  CHECK(context.is_equal(p));
-
+std::vector<long> Scop::getParameterValues() const {
   // Scop holds a vector of Variables.
   // Iterate over parameters in order, checking if the
-  // context contains a parameter whose name corresponds to that
+  // map of known parameter values contains a parameter
+  // whose name corresponds to that
   // Variable and push respective parameter values.
   std::vector<long> paramValues;
   for (auto const& param : halide.params) {
-    isl::id id(ctx, param.name());
-    CHECK(context.involves_param(id));
-    auto val = isl::aff::param_on_domain_space(space, id).eval(p);
-    CHECK(val.is_int()) << "fractional parameters unsupported";
-    CHECK(val.le(longMax)) << "parameter value overflows long";
-    paramValues.push_back(val.get_num_si());
+    auto name = param.name();
+    CHECK(parameterValues.count(name) == 1);
+    paramValues.push_back(parameterValues.at(name));
   }
   return paramValues;
 }
diff --git a/tc/core/polyhedral/scop.h b/tc/core/polyhedral/scop.h
diff --git a/test/test_cuda_mapper.cc b/test/test_cuda_mapper.cc
diff --git a/test/test_cuda_mapper_memory_promotion.cc b/test/test_cuda_mapper_memory_promotion.cc
diff --git a/test/test_mapper_llvm.cc b/test/test_mapper_llvm.cc
diff --git a/test/test_mapper_tapir.cc b/test/test_mapper_tapir.cc

Original file line number	Diff line number	Diff line change
`@@ -227,19 +227,19 @@ isl::aff makeIslAffFromExpr(isl::space space, const Expr& e) {`
`227`	`227`	`return list[0];`
`228`	`228`	`}`
`229`	`229`
`230`		`-isl::space makeParamSpace(isl::ctx ctx, const SymbolTable& symbolTable) {`
	`230`	`+isl::space makeParamSpace(isl::ctx ctx, const ParameterVector& params) {`
`231`	`231`	`auto space = isl::space(ctx, 0);`
`232`	`232`	`// set parameter names`
`233`		`- for (auto p : symbolTable.params) {`
	`233`	`+ for (auto p : params) {`
`234`	`234`	`space = space.add_param(isl::id(ctx, p.name()));`
`235`	`235`	`}`
`236`	`236`	`return space;`
`237`	`237`	`}`
`238`	`238`
`239`		`-isl::set makeParamContext(isl::ctx ctx, const SymbolTable& symbolTable) {`
`240`		`- auto space = makeParamSpace(ctx, symbolTable);`
	`239`	`+isl::set makeParamContext(isl::ctx ctx, const ParameterVector& params) {`
	`240`	`+ auto space = makeParamSpace(ctx, params);`
`241`	`241`	`auto context = isl::set::universe(space);`
`242`		`- for (auto p : symbolTable.params) {`
	`242`	`+ for (auto p : params) {`
`243`	`243`	`isl::aff a(isl::aff::param_on_domain_space(space, isl::id(ctx, p.name())));`
`244`	`244`	`context = context & (a >= 0);`
`245`	`245`	`}`
Original file line number	Diff line number	Diff line change
`@@ -421,10 +421,7 @@ isl::set tensorElementsSet(const Scop& scop, isl::id tensorId) {`
`421`	`421`	`(isl::aff_set(aff) < (minAff + extentAff));`
`422`	`422`	`}`
`423`	`423`
`424`		`- if (scop.globalParameterContext) {`
`425`		`- tensorElements =`
`426`		`- tensorElements.intersect_params(scop.globalParameterContext);`
`427`		`- }`
	`424`	`+ tensorElements = tensorElements.intersect_params(scop.context());`
`428`	`425`	`return tensorElements;`
`429`	`426`	`}`
`430`	`427`