[consan] Handle all tmem allocations (#8787)

ThomasRaoux · anmyachev · commit 20095e601fa4 · 2025-12-01T10:59:17.000Z
diff --git a/lib/Conversion/TritonInstrumentToLLVM/InstrumentationToLLVM.cpp b/lib/Conversion/TritonInstrumentToLLVM/InstrumentationToLLVM.cpp
@@ -27,7 +27,7 @@ Value createMemDescToI64(RewriterBase &rewriter, Location loc,
                          const LLVMTypeConverter *typeConverter,
                          ttg::MemDescType memDescTy, Value sharedMemStruct) {
   TritonLLVMOpBuilder b(loc, rewriter);
-  if (isa<ttng::TensorMemoryEncodingAttr>(memDescTy.getEncoding())) {
+  if (isa<ttng::TensorMemorySpaceAttr>(memDescTy.getMemorySpace())) {
     return b.ptrtoint(rewriter.getIntegerType(64), sharedMemStruct);
   }
   assert(isa<ttg::SharedEncodingTrait>(memDescTy.getEncoding()) &&

Original file line number	Diff line number	Diff line change
`@@ -27,7 +27,7 @@ Value createMemDescToI64(RewriterBase &rewriter, Location loc,`
`27`	`27`	`const LLVMTypeConverter *typeConverter,`
`28`	`28`	`ttg::MemDescType memDescTy, Value sharedMemStruct) {`
`29`	`29`	`TritonLLVMOpBuilder b(loc, rewriter);`
`30`		`- if (isa<ttng::TensorMemoryEncodingAttr>(memDescTy.getEncoding())) {`
	`30`	`+ if (isa<ttng::TensorMemorySpaceAttr>(memDescTy.getMemorySpace())) {`
`31`	`31`	`return b.ptrtoint(rewriter.getIntegerType(64), sharedMemStruct);`
`32`	`32`	`}`
`33`	`33`	`assert(isa<ttg::SharedEncodingTrait>(memDescTy.getEncoding()) &&`