[Feat] Add tilelang T.assume support and assume injection for buffer shapes

src/transform/inject_assumes.cc

@@ -0,0 +1,108 @@
+
+#include "tvm/arith/analyzer.h"
+#include "tvm/ir/expr.h"
+#include "tvm/ir/transform.h"
+#include "tvm/node/structural_hash.h"
+#include "tvm/tir/expr.h"
+#include "tvm/tir/stmt.h"
+#include "tvm/tir/stmt_functor.h"
+#include "tvm/tir/transform.h"
+
+namespace tvm::tl {
+using namespace tir;
+
+class AssumeInjector : public tvm::tir::StmtExprMutator {
+  using Base = tvm::tir::StmtExprMutator;
+public:
+  AssumeInjector(PrimFunc f): f(f) {}
+  static PrimFunc Substitute(PrimFunc f) {
+    auto injector = AssumeInjector(f);
+    f.CopyOnWrite()->body = injector(f->body);
+    return f;
+  }
+private:
+  struct AssertCreator {
+    tvm::StructuralHash sh;
+    tvm::StructuralEqual se;
+    std::unordered_map<size_t, std::vector<tvm::PrimExpr>> buckets;
+    std::vector<PrimExpr> exprs;
+    void addExpr(PrimExpr e) {
+      size_t h = sh(e);
+      auto bucket = buckets[h];
+      auto it = std::find_if(bucket.begin(), bucket.end(), [&](auto y) {
+        return se(e, y, true);
+      });
+      if(it == bucket.end()) {
+        exprs.push_back(e);
+        buckets[h].push_back(e);
+      }
+    }
+    void addBuffer(Buffer buf) {
+      for(auto shape: buf->shape) {
+        if(shape->IsInstance<IntImmNode>()) continue;
+        addExpr(shape);
+      }
+    }
+    Stmt build(Stmt body) {
+      if(exprs.empty()) return body;
+      PrimExpr red = GT(exprs[0], 0);
+      for(size_t i = 1; i < exprs.size(); ++i) {
+        red = And(GT(exprs[i], 0), red);
+      }
+      auto simplified = arith::Analyzer{}.Simplify(red, 10);
+      auto msg = StringImm("Invalid Buffer Shape: buffer shape should be greater than 0");
+      return AssertStmt(simplified, msg, body);
+    }
+  };
+  Stmt VisitStmt_(const DeclBufferNode* op) final {
+    auto body = VisitStmt(op->body);
+    AssertCreator c;
+    c.addBuffer(op->buffer);
+    return DeclBuffer(op->buffer, c.build(body), op->span);
+  }
+  Stmt VisitStmt_(const BlockNode* op) final {
+    auto body = VisitStmt(op->body);
+    AssertCreator c;
+    if(root_node) {
+      for(auto item: f->buffer_map) {
+        c.addBuffer(item.second);
+      }
+    }
+    for(auto item: op->alloc_buffers) {
+      c.addBuffer(item);
+    }
+    for(auto item: op->match_buffers) {
+      c.addBuffer(item->buffer);
+    }
+    return Block(
+      op->iter_vars,
+      op->reads,
+      op->writes,
+      op->name_hint,
+      c.build(body),
+      op->init,
+      op->alloc_buffers,
+      op->match_buffers,
+      op->annotations,
+      op->span
+    );
+  }
+  PrimFunc f;
+  bool root_node { true };
+};
+
+using namespace tir::transform;
+
+tvm::transform::Pass InjectAssumes() {
+  auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
+    return AssumeInjector::Substitute(f);
+  };
+  return CreatePrimFuncPass(pass_func, 0, "tl.InjectAssumes", {});
+}
+
+TVM_FFI_STATIC_INIT_BLOCK({
+  namespace refl = tvm::ffi::reflection;
+  refl::GlobalDef().def("tl.transform.InjectAssumes", InjectAssumes);
+});
+
+} // namespace tvm::tl

tilelang/engine/phase.py

@@ -4,6 +4,7 @@
 from tilelang.transform import PassContext
 from tilelang.contrib.nvcc import have_tma
 from typing import Optional
+from pathlib import Path
 
 
 def allow_warp_specialized(pass_ctx: Optional[PassContext] = None,
@@ -61,6 +62,10 @@ def should_enable_aggressive_merge(pass_ctx: Optional[PassContext] = None,
     return enable_aggressive_merge
 
 
+debug_path = Path('debug')
+debug_path.mkdir(exist_ok=True)
+
+
 def LowerAndLegalize(mod: IRModule, target: Target) -> IRModule:
     # Bind the target device information to the module
     """
@@ -87,14 +92,18 @@ def LowerAndLegalize(mod: IRModule, target: Target) -> IRModule:
 
     # Legalize the frontend IR to make it compatible with TVM
     mod = tilelang.transform.FrontendLegalize()(mod)
+    # Inject assumes to speedup tvm prover
+    mod = tilelang.transform.InjectAssumes()(mod)
     # Simplify the IR expressions
     mod = tir.transform.Simplify()(mod)
     # Set layouts for reducers
     mod = tilelang.transform.LayoutReducer()(mod)
     # Infer memory layouts for fragments and shared memory
     mod = tilelang.transform.LayoutInference()(mod)
+    debug_path.joinpath('LowerTileOp.0.py').write_text(mod.script(show_meta=True))
     # Lower high-level tile operations to low-level operations
     mod = tilelang.transform.LowerTileOp()(mod)
+    debug_path.joinpath('LowerTileOp.1.py').write_text(mod.script(show_meta=True))
     # Lower l2 persistent map
     mod = tilelang.transform.LowerL2Persistent()(mod)
     # Legalize vectorized loops to ensure they are valid
@@ -133,7 +142,9 @@ def OptimizeForTarget(mod: IRModule, target: Target) -> IRModule:
         mod = tilelang.transform.LowerOpaqueBlock()(mod)
         mod = tilelang.transform.MergeIfStmt()(mod)
         mod = tilelang.transform.RewriteWgmmaSync()(mod)
+        debug_path.joinpath('InjectFenceProxy.0.py').write_text(mod.script(show_meta=True))
         mod = tilelang.transform.InjectFenceProxy()(mod)
+        debug_path.joinpath('InjectFenceProxy.1.py').write_text(mod.script(show_meta=True))
     else:
         mod = tilelang.transform.IfStmtBinding()(mod)
         mod = tir.transform.PlanAndUpdateBufferAllocationLocation()(mod)
@@ -197,5 +208,6 @@ def OptimizeForTarget(mod: IRModule, target: Target) -> IRModule:
 
     # Transform threadblock to persistent threadblock
     mod = tilelang.transform.PersistThreadblock()(mod)
+    debug_path.joinpath('PersistThreadblock.1.py').write_text(mod.script(show_meta=True))
 
     return mod

tilelang/transform/__init__.py

@@ -78,6 +78,16 @@ def FrontendLegalize():
     """
     return _ffi_api.FrontendLegalize()  # type: ignore
 
+def InjectAssumes():
+    """Inject Assumes
+
+    Returns:
+    -------
+    fpass : tvm.transform.Pass
+        The result pass
+    """
+    return _ffi_api.InjectAssumes()
+
 
 def LowerHopperIntrin():
     """LowerHopperIntrin