Merge pull request #578 from facebookresearch/pr/pre-template

prepare for templated isl types

Author: skimo-openhub

tc/core/CMakeLists.txt

@@ -34,6 +34,7 @@ add_library(
   polyhedral/scop.cc
   polyhedral/separation.cc
   polyhedral/unroll.cc
+  polyhedral/utils.cc
 )
 target_include_directories(tc_core PUBLIC ${LLVM_INCLUDE_DIRS})
 target_link_libraries(

tc/core/halide2isl.cc

@@ -24,6 +24,7 @@
 #include "tc/core/polyhedral/schedule_isl_conversion.h"
 #include "tc/core/polyhedral/schedule_transforms.h"
 #include "tc/core/polyhedral/schedule_tree.h"
+#include "tc/core/polyhedral/utils.h"
 #include "tc/core/tc2halide.h"
 
 namespace tc {
@@ -259,7 +260,8 @@ isl::map extractAccess(
 
   isl::space paramSpace = domain.paramSpace;
   isl::id tensorID(paramSpace.get_ctx(), tensor);
-  auto tensorSpace = paramSpace.add_named_tuple_id_ui(tensorID, args.size());
+  auto tensorTuple = constructTensorTuple(paramSpace, tensorID, args.size());
+  auto tensorSpace = tensorTuple.get_space();
 
   // Start with a totally unconstrained set - every point in
   // the allocation could be accessed.
@@ -275,9 +277,10 @@ isl::map extractAccess(
     // The coordinate written to in the range ...
     auto rangePoint = identity.get_aff(i);
     // ... equals the coordinate accessed as a function of the parameters.
-    auto domainPoint = halide2isl::makeIslAffFromExpr(tensorSpace, args[i]);
+    auto domainPoint = halide2isl::makeIslAffFromExpr(paramSpace, args[i]);
     if (!domainPoint.is_null()) {
-      access = access.intersect(isl::pw_aff(domainPoint).eq_set(rangePoint));
+      domainPoint = domainPoint.unbind_params_insert_domain(tensorTuple);
+      access = access.intersect(domainPoint.eq_set(rangePoint));
     }
   }
 

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -573,7 +573,7 @@ Scop::SyncLevel MappedScop::findBestSync(
   auto contextSt = scop_->scheduleRoot()->children()[0];
   auto contextElem = contextSt->as<detail::ScheduleTreeContext>();
   TC_CHECK(nullptr != contextElem);
-  dependences = dependences.intersect_params(contextElem->context_);
+  dependences = dependences.intersect_params(contextElem->context_.params());
 
   if (dependences.is_subset(dependences.eq_at(domainToThread))) {
     return Scop::SyncLevel::None;

tc/core/polyhedral/memory_promotion.cc

@@ -25,6 +25,7 @@
 #include "tc/core/polyhedral/exceptions.h"
 #include "tc/core/polyhedral/schedule_tree.h"
 #include "tc/core/polyhedral/scop.h"
+#include "tc/core/polyhedral/utils.h"
 #include "tc/external/isl.h"
 
 namespace tc {
@@ -409,17 +410,20 @@ namespace {
 // context of the scop.
 isl::set tensorElementsSet(const Scop& scop, isl::id tensorId) {
   auto halideParameter = scop.findArgument(tensorId).parameter();
-  auto space = scop.domain().get_space().params();
+  auto space = scop.domain().get_space();
   auto nDim = halideParameter.dimensions();
-  space = space.add_named_tuple_id_ui(tensorId, nDim);
+  auto tensorTuple = constructTensorTuple(space, tensorId, nDim);
+  auto tensorSpace = tensorTuple.get_space();
 
-  auto tensorElements = isl::set::universe(space);
-  auto identity = isl::multi_aff::identity(space.range().map_from_set());
+  auto tensorElements = isl::set::universe(tensorSpace);
+  auto identity = isl::multi_aff::identity(tensorSpace.map_from_set());
   for (int i = 0; i < nDim; ++i) {
     auto minAff = halide2isl::makeIslAffFromExpr(
         space, halideParameter.min_constraint(i));
     auto extentAff = halide2isl::makeIslAffFromExpr(
         space, halideParameter.extent_constraint(i));
+    minAff = minAff.unbind_params_insert_domain(tensorTuple);
+    extentAff = extentAff.unbind_params_insert_domain(tensorTuple);
     auto aff = identity.get_aff(i);
     tensorElements = tensorElements & (minAff <= isl::aff_set(aff)) &
         (isl::aff_set(aff) < (minAff + extentAff));

tc/core/polyhedral/scop.h

@@ -136,7 +136,7 @@ struct Scop {
   isl::set makeContext(
       const std::unordered_map<std::string, T>& sizes =
           std::unordered_map<std::string, T>()) const {
-    auto s = domain().get_space().params();
+    auto s = domain().get_space();
     return makeSpecializationSet(s, sizes);
   }
 

tc/core/polyhedral/utils.cc

@@ -0,0 +1,39 @@
+/**
+ * Copyright (c) 2018, Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "tc/core/polyhedral/utils.h"
+
+namespace tc {
+namespace polyhedral {
+
+/* Construct a tuple representing the tensor with identifier "tensorId" and
+ * dimension "dim" from the parameter space "paramSpace",
+ * without any specific names for the indices, from the perspective
+ * of the user.
+ * Since some names are required, use names of the form "__tc_tensor_arg*".
+ */
+isl::multi_id
+constructTensorTuple(isl::space paramSpace, isl::id tensorId, size_t dim) {
+  auto tensorSpace = paramSpace.add_named_tuple_id_ui(tensorId, dim);
+  isl::id_list tensorArgs(paramSpace.get_ctx(), 0);
+  for (size_t i = 0; i < dim; ++i) {
+    auto name = std::string("__tc_tensor_arg") + std::to_string(i);
+    tensorArgs = tensorArgs.add(isl::id(paramSpace.get_ctx(), name));
+  }
+  return isl::multi_id(tensorSpace, tensorArgs);
+}
+
+} // namespace polyhedral
+} // namespace tc

tc/core/polyhedral/utils.h

@@ -0,0 +1,31 @@
+/**
+ * Copyright (c) 2018, Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include "tc/external/isl.h"
+
+namespace tc {
+namespace polyhedral {
+
+/* Construct a tuple representing the tensor with identifier "tensorId" and
+ * dimension "dim" from the parameter space "paramSpace",
+ * without any specific names for the indices.
+ */
+isl::multi_id
+constructTensorTuple(isl::space paramSpace, isl::id tensorId, size_t dim);
+
+} // namespace polyhedral
+} // namespace tc

tc/external/detail/islpp-inl.h

@@ -41,9 +41,9 @@ inline isl::aff operator/(isl::aff A, int i) {
   return A.scale_down(isl::val(A.get_ctx(), i));
 }
 
-inline isl::aff operator+(int i, isl::aff A) {
-  isl::ctx ctx = A.get_ctx();
-  return A + isl::val(ctx, i);
+template <typename T>
+inline isl::aff operator+(int i, T A) {
+  return A.add_constant_si(i);
 }
 
 inline isl::aff operator+(isl::aff A, isl::val v) {
@@ -55,10 +55,6 @@ inline isl::aff operator+(isl::val v, isl::aff A) {
   return A + v;
 }
 
-inline isl::aff operator+(isl::aff A, isl::aff B) {
-  return A.add(B);
-}
-
 inline isl::aff operator+(isl::aff A, int i) {
   return i + A;
 }
@@ -195,10 +191,6 @@ inline isl::multi_aff operator/(isl::multi_aff left, isl::multi_val right) {
 ///////////////////////////////////////////////////////////////////////////////
 // Operations on isl::set and isl::union_set
 ///////////////////////////////////////////////////////////////////////////////
-inline isl::set operator&(isl::set S1, isl::set S2) {
-  return S1.intersect(S2);
-}
-
 inline isl::union_set operator&(isl::union_set S1, isl::set S2) {
   return S1.intersect(isl::union_set(S2));
 }
@@ -207,10 +199,6 @@ inline isl::union_set operator&(isl::set S1, isl::union_set S2) {
   return S2 & S1;
 }
 
-inline isl::union_set operator&(isl::union_set S1, isl::union_set S2) {
-  return S1.intersect(S2);
-}
-
 ///////////////////////////////////////////////////////////////////////////////
 // Operations on isl::set and isl::point
 ///////////////////////////////////////////////////////////////////////////////

tc/external/detail/islpp.h

@@ -35,11 +35,21 @@ namespace isl {
 // the official ISL C++ bindings.
 //
 
+template <typename T>
+inline T operator+(T a, T b) {
+  return a.add(b);
+}
+
 template <typename T>
 inline T operator-(T a, T b) {
   return a.sub(b);
 }
 
+template <typename T>
+inline T operator&(T S1, T S2) {
+  return S1.intersect(S2);
+}
+
 inline isl::val operator*(isl::val l, isl::val r) {
   return l.mul(r);
 }
@@ -52,10 +62,6 @@ inline isl::val operator*(long i, isl::val v) {
   return v * i;
 }
 
-inline isl::val operator+(isl::val l, isl::val r) {
-  return l.add(r);
-}
-
 inline isl::val operator+(isl::val v, long i) {
   return v.add(isl::val(v.get_ctx(), i));
 }
@@ -122,8 +128,6 @@ isl::aff operator*(isl::val V, isl::aff A);
 
 isl::aff operator/(isl::aff A, int i);
 
-isl::aff operator+(int i, isl::aff A);
-isl::aff operator+(isl::aff A, isl::aff B);
 isl::aff operator+(isl::aff A, int i);
 isl::aff operator+(isl::aff A, isl::val v);
 isl::aff operator+(isl::val v, isl::aff A);
@@ -184,10 +188,8 @@ isl::multi_aff operator/(isl::multi_aff left, isl::multi_val right);
 ///////////////////////////////////////////////////////////////////////////////
 // Operations on isl::set and isl::union_set
 ///////////////////////////////////////////////////////////////////////////////
-isl::set operator&(isl::set S1, isl::set S2);
 isl::union_set operator&(isl::union_set S1, isl::set S2);
 isl::union_set operator&(isl::set S1, isl::union_set S2);
-isl::union_set operator&(isl::union_set S1, isl::union_set S2);
 
 ///////////////////////////////////////////////////////////////////////////////
 // Operations on isl::set and isl::point