Add a stripmineSink and imperfectly nested tiling primitives.
This CL adds a primitive to perform stripmining of a loop by a given factor and sinking it under multiple target loops. In turn this is used to implement imperfectly nested loop tiling (with interchange) by repeatedly calling the stripmineSink primitive. The API returns the point loops and allows repeated invocations of tiling to achieve declarative, multi-level, imperfectly-nested tiling. Note that this CL is only concerned with the mechanical aspects and does not worry about analysis and legality. The API is demonstrated in an example which creates an EDSC block, emits the corresponding MLIR and applies imperfectly-nested tiling: ```cpp auto block = edsc::block({ For(ArrayRef<edsc::Expr>{i, j}, {zero, zero}, {M, N}, {one, one}, { For(k1, zero, O, one, { C({i, j, k1}) = A({i, j, k1}) + B({i, j, k1}) }), For(k2, zero, O, one, { C({i, j, k2}) = A({i, j, k2}) + B({i, j, k2}) }), }), }); // clang-format on emitter.emitStmts(block.getBody()); auto l_i = emitter.getAffineForOp(i), l_j = emitter.getAffineForOp(j), l_k1 = emitter.getAffineForOp(k1), l_k2 = emitter.getAffineForOp(k2); auto indicesL1 = mlir::tile({l_i, l_j}, {512, 1024}, {l_k1, l_k2}); auto l_ii1 = indicesL1[0][0], l_jj1 = indicesL1[1][0]; mlir::tile({l_jj1, l_ii1}, {32, 16}, l_jj1); ``` The edsc::Expr for the induction variables (i, j, k_1, k_2) provide the programmatic hooks from which tiling can be applied declaratively. PiperOrigin-RevId: 235548228
Loading
Please sign in to comment