Use static shapes (#25)

* generate static shapes at function boundary
* wave equation: use colon slice where possible

Co-authored-by: Frank Schlimbach <frank.schlimbach@intel.com>

Author: tkarna

examples/wave_equation.py

@@ -144,7 +144,7 @@ def exact_elev(t, x_t_2d, y_t_2d, lx, ly):
         return amp * sol_x * sol_y * sol_t
 
     # inital elevation
-    e[0:nx, 0:ny] = exact_elev(0.0, x_t_2d, y_t_2d, lx, ly)
+    e[:, :] = exact_elev(0.0, x_t_2d, y_t_2d, lx, ly)
 
     # compute time step
     alpha = 0.5
@@ -173,13 +173,12 @@ def rhs(u, v, e):
         # sign convention: positive on rhs
 
         # pressure gradient -g grad(elev)
-        dudt = -g * (e[1:nx, 0:ny] - e[0 : nx - 1, 0:ny]) / dx
-        dvdt = -g * (e[0:nx, 1:ny] - e[0:nx, 0 : ny - 1]) / dy
+        dudt = -g * (e[1:nx, :] - e[0 : nx - 1, :]) / dx
+        dvdt = -g * (e[:, 1:ny] - e[:, 0 : ny - 1]) / dy
 
         # velocity divergence -h div(u)
         dedt = -h * (
-            (u[1 : nx + 1, 0:ny] - u[0:nx, 0:ny]) / dx
-            + (v[0:nx, 1 : ny + 1] - v[0:nx, 0:ny]) / dy
+            (u[1 : nx + 1, :] - u[0:nx, :]) / dx + (v[:, 1 : ny + 1] - v[:, 0:ny]) / dy
         )
 
         return dudt, dvdt, dedt
@@ -189,19 +188,19 @@ def step(u, v, e, u1, v1, e1, u2, v2, e2):
         Execute one SSPRK(3,3) time step
         """
         dudt, dvdt, dedt = rhs(u, v, e)
-        u1[1:nx, 0:ny] = u[1:nx, 0:ny] + dt * dudt
-        v1[0:nx, 1:ny] = v[0:nx, 1:ny] + dt * dvdt
-        e1[0:nx, 0:ny] = e[0:nx, 0:ny] + dt * dedt
+        u1[1:nx, :] = u[1:nx, :] + dt * dudt
+        v1[:, 1:ny] = v[:, 1:ny] + dt * dvdt
+        e1[:, :] = e[:, :] + dt * dedt
 
         dudt, dvdt, dedt = rhs(u1, v1, e1)
-        u2[1:nx, 0:ny] = 0.75 * u[1:nx, 0:ny] + 0.25 * (u1[1:nx, 0:ny] + dt * dudt)
-        v2[0:nx, 1:ny] = 0.75 * v[0:nx, 1:ny] + 0.25 * (v1[0:nx, 1:ny] + dt * dvdt)
-        e2[0:nx, 0:ny] = 0.75 * e[0:nx, 0:ny] + 0.25 * (e1[0:nx, 0:ny] + dt * dedt)
+        u2[1:nx, :] = 0.75 * u[1:nx, :] + 0.25 * (u1[1:nx, :] + dt * dudt)
+        v2[:, 1:ny] = 0.75 * v[:, 1:ny] + 0.25 * (v1[:, 1:ny] + dt * dvdt)
+        e2[:, :] = 0.75 * e[:, :] + 0.25 * (e1[:, :] + dt * dedt)
 
         dudt, dvdt, dedt = rhs(u2, v2, e2)
-        u[1:nx, 0:ny] = u[1:nx, 0:ny] / 3.0 + 2.0 / 3.0 * (u2[1:nx, 0:ny] + dt * dudt)
-        v[0:nx, 1:ny] = v[0:nx, 1:ny] / 3.0 + 2.0 / 3.0 * (v2[0:nx, 1:ny] + dt * dvdt)
-        e[0:nx, 0:ny] = e[0:nx, 0:ny] / 3.0 + 2.0 / 3.0 * (e2[0:nx, 0:ny] + dt * dedt)
+        u[1:nx, :] = u[1:nx, :] / 3.0 + 2.0 / 3.0 * (u2[1:nx, :] + dt * dudt)
+        v[:, 1:ny] = v[:, 1:ny] / 3.0 + 2.0 / 3.0 * (v2[:, 1:ny] + dt * dvdt)
+        e[:, :] = e[:, :] / 3.0 + 2.0 / 3.0 * (e2[:, :] + dt * dedt)
 
     t = 0
     i_export = 0

imex_version.txt

@@ -1 +1 @@
-d944080e43cc67798b92f1ca5d1ea8e50df69eb5
+2726c17193e3297a5442ba507ba560133723ac09

src/EWBinOp.cpp

@@ -103,6 +103,9 @@ struct DeferredEWBinOp : public Deferred {
 
     auto bop = builder.create<::imex::ptensor::EWBinOp>(
         loc, outTyp, builder.getI32IntegerAttr(ddpt2mlir(_op)), av, bv);
+    // auto bop =
+    //     builder.create<::imex::ptensor::EWBinOp>(loc, ddpt2mlir(_op), av,
+    //     bv);
     dm.addVal(this->guid(), bop,
               [this](Transceiver *transceiver, uint64_t rank, void *allocated,
                      void *aligned, intptr_t offset, const intptr_t *sizes,

src/jit/mlir.cpp

@@ -34,6 +34,7 @@
 */
 
 #include "ddptensor/jit/mlir.hpp"
+#include "ddptensor/DDPTensorImpl.hpp"
 #include "ddptensor/Registry.hpp"
 
 #include "mlir/IR/MLIRContext.h"
@@ -116,7 +117,8 @@ static ::mlir::Type makeSignlessType(::mlir::Type type) {
 
 // convert ddpt's DTYpeId into MLIR type
 static ::mlir::Type getTType(::mlir::OpBuilder &builder, DTypeId dtype,
-                             int rank, uint64_t team, bool balanced) {
+                             ::mlir::SmallVector<int64_t> &shape, uint64_t team,
+                             bool balanced) {
   ::mlir::Type etyp;
 
   switch (dtype) {
@@ -149,7 +151,6 @@ static ::mlir::Type getTType(::mlir::OpBuilder &builder, DTypeId dtype,
     throw std::runtime_error("unknown dtype");
   };
 
-  ::mlir::SmallVector<int64_t> shape(rank, ::mlir::ShapedType::kDynamic);
   auto pttype = ::imex::ptensor::PTensorType::get(shape, etyp);
   if (team) {
     return ::imex::dist::DistTensorType::get(builder.getContext(), pttype);
@@ -165,8 +166,15 @@ ::mlir::Value DepManager::getDependent(::mlir::OpBuilder &builder,
     // Not found -> this must be an input argument to the jit function
     auto idx = _args.size();
     auto fut = Registry::get(guid);
+    auto impl = std::dynamic_pointer_cast<DDPTensorImpl>(fut.get());
+    auto rank = impl->ndims();
+    const int64_t *shape_ptr = impl->local_shape();
+    ::mlir::SmallVector<int64_t> shape(rank);
+    for (size_t i = 0; i < rank; i++) {
+      shape[i] = shape_ptr[i];
+    }
     auto typ =
-        getTType(builder, fut.dtype(), fut.rank(), fut.team(), fut.balanced());
+        getTType(builder, fut.dtype(), shape, fut.team(), fut.balanced());
     _func.insertArgument(idx, typ, {}, loc);
     auto val = _func.getArgument(idx);
     _args.push_back({guid, std::move(fut)});