Skip to content

Init identity matrix and jump to main #611

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ppotapov-aws merged 1 commit into from
Mar 6, 2026
Merged

Init identity matrix and jump to main #611

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 18 additions & 2 deletions KLR/Trace/Types.lean
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -429,7 +429,13 @@ def addId : Trace Unit := do
throw "identity already initialized"
let dtype := .int8
let shape := Core.Shape.mk 128 [128]
let tensorName <- Core.TensorName.make idName.toString dtype shape none (<- flags.address_rotation)
let idSbufAddr : Address := {
name := idName.toString,
memory := .sbuf,
parSize := 128
freeSize := 128
}
let tensorName <- Core.TensorName.make idName.toString dtype shape idSbufAddr (<- flags.address_rotation)
let id : KLR.Core.TensorRef := .abstract (.simple tensorName)
let pos : Pos := { line := 0, column := 0 }
let hbmInitName := <-genName
Expand All @@ -452,8 +458,18 @@ def addId : Trace Unit := do
}) none pos
let lbl := (<- genLabel `init)
let idTensor := identity 128
-- Jump from init block to the first main block so the CFG has a single exit.
Comment thread
ppotapov-aws marked this conversation as resolved.
let mainLabel := match (<- get).label with
| some l => l
| none => "main.1" -- fallback; should not happen
let jmpName := (<- genLabel `jmp)
let jmpStmt := Core.Stmt.oper (.cmpBranch {
reg1 := "", reg2 := "", imm := 0,
op := BrCmpOp.always,
trueLabel := mainLabel, falseLabel := ""
}) jmpName pos
modify fun s => { s with
body := #[Block.mk lbl false [initStmt]] ++ s.body,
body := #[Block.mk lbl false [initStmt, jmpStmt]] ++ s.body,
sharedConstants := s.sharedConstants.push (hbmInitName.toString, idTensor)
}
extend_global idName (.access (.simple tensorName))
Expand Down
Loading