Pykernel - matmul common reader/writer kernels #2289

vtangTT · 2025-02-26T03:26:16Z

Ticket

#2254

More specifically:
#2178
#2247
#2249

Problem description

Want to be able to write reader and writer kernels used in matmul programming examples.

What's changed

InterleavedAddrGenFast cpp struct mapping:

struct itself is exposed to python through TTKernelOp get_interleaved_addr_gen_fast
conversionpatternrewriter for said Op that will do the following:

%6 = "ttkernel.get_interleaved_addr_gen_fast"(%3, %0, %4, %5) : (i1, i32, i32, !ttkernel.DataFormat) -> !ttkernel.interleaved_addr_gen_fast

-> to emitc

%14 = "emitc.variable"() <{value = #emitc.opaque<"">}> : () -> !emitc.lvalue<!emitc.opaque<"InterleavedAddrGenFast<true>">>
%15 = "emitc.member"(%14) <{member = "bank_base_address"}> : (!emitc.lvalue<!emitc.opaque<"InterleavedAddrGenFast<true>">>) -> !emitc.lvalue<i32>
%16 = "emitc.member"(%14) <{member = "page_size"}> : (!emitc.lvalue<!emitc.opaque<"InterleavedAddrGenFast<true>">>) -> !emitc.lvalue<i32>
%17 = "emitc.member"(%14) <{member = "data_format"}> : (!emitc.lvalue<!emitc.opaque<"InterleavedAddrGenFast<true>">>) -> !emitc.lvalue<!emitc.opaque<"DataFormat">>
emitc.assign %5 : i32 to %15 : <i32>
emitc.assign %12 : i32 to %16 : <i32>
emitc.assign %13 : !emitc.opaque<"DataFormat"> to %17 : <!emitc.opaque<"DataFormat">>
%18 = emitc.load %14 : <!emitc.opaque<"InterleavedAddrGenFast<true>">>

-> to cpp

InterleavedAddrGenFast<true> v15;
v15.bank_base_address = v6;
v15.page_size = v13;
v15.data_format = v14;
InterleavedAddrGenFast<true> v16 = v15;

added noc_async_read_tile ttkernel op
added noc_async_write_tile ttkernel op

Checklist

TTKernelToEmitC unit tests for all added ops

vroubtsovTT · 2025-02-27T02:15:01Z

include/ttmlir/Dialect/TTKernel/IR/TTKernelOps.td

+      NocAsyncReadTile
+    }];
+
+    let arguments = (ins I32:$id, TTKernel_InterleavedAddrGenFast:$s, I32:$dstLocalL1Addr);


s is too short of an identifier to be useful, can we change it to something longer/more self-documenting?

changed to addrGenStruct

vroubtsovTT · 2025-02-27T02:16:45Z

include/ttmlir/Dialect/TTKernel/IR/TTKernelOps.td

@@ -634,6 +652,16 @@ def TTKernel_StoreToL1Op : TTKernel_Op<"store_to_l1"> {
    let arguments = (ins I32:$value, TTKernel_L1AddrPtr:$l1_ptr, I32:$offset);
 }

+def TTKernel_GetInterleavedAddrGenFastOp : TTKernel_Op<"get_interleaved_addr_gen_fast"> {
+    let summary = "GetAddrGenFastConfig";


summary doesn't seem to match GetInterleavedAddrGenFastOp

vroubtsovTT · 2025-02-27T02:20:17Z

lib/Conversion/TTKernelToEmitC/TTKernelToEmitC.cpp

+          rewriter.create<emitc::LoadOp>(op->getLoc(), opaqueStructType, varOp);
+
+      // Replace the original operation with the loaded value so it can be used.
+      op.replaceAllUsesWith(loadOp.getResult());


I don't think op.replaceAllUsesWith(loadOp.getResult()); is a legal modification when using conversion patterns -- all IR modifications must be done via rewriter (see https://www.youtube.com/watch?v=xIeihq2WZOU). The rewriter.replaceOp... line should be enough -- is that not the case?

Yes, I'm not sure why I left that in there. It works without, thank you!

vroubtsovTT · 2025-02-27T02:21:46Z

include/ttmlir/Dialect/TTKernel/IR/TTKernelOps.td

+      NocAsyncWriteTilie
+    }];
+
+    let arguments = (ins I32:$id, TTKernel_InterleavedAddrGenFast:$s, I32:$srcLocalL1Addr);


Same comment about s.

changed to addrGenStruct

vroubtsovTT · 2025-02-27T02:24:42Z

lib/Conversion/TTKernelToEmitC/TTKernelToEmitC.cpp

@@ -556,6 +624,7 @@ class ConvertTTKernelToEmitCPass
        return op.getNumArguments() == 0;
      });
      target.addLegalOp<func::ReturnOp>();
+      target.addIllegalOp<ttkernel::GetInterleavedAddrGenFastOp>();


this should not be needed here because the next line declares all of ttkernel dialect as illegal at the end of this conversion step

Yep, this was sloppy, ty!

vroubtsovTT · 2025-02-27T02:30:00Z

test/ttmlir/Conversion/TTKernelToEmitC/ttkernel.mlir

+      // CHECK: "emitc.member"(%[[VAR]]) <{member = "data_format"}>{{.*}}
+      // CHECK: emitc.assign {{.*}}
+      // CHECK: emitc.assign {{.*}}
+      // CHECK: emitc.assign {{.*}}


using {{.*}} at the end of regexes isn't doing anything useful, because CHECK is already a match for a substring of a line

removed ty!

vroubtsovTT · 2025-02-27T02:38:02Z

test/ttmlir/Conversion/TTKernelToEmitC/ttkernel.mlir

+      // CHECK: = emitc.load %[[VAR]] : <!emitc.opaque<"InterleavedAddrGenFast<true>">>
+      %is_dram = arith.constant 1 : i1
+      %temp1 = arith.constant 262400 : i32
+      %temp2 = arith.constant 32 : i32


asking for two changes here:

since we are using %temp1 and %temp2 in multiple places later and they aren't just dummy inputs, can we give these SSA vars some meaningful names?

can we capture the names of these variables by using %[[CAPITALIZED_SSA_VAR_NAME]] = ... CHECKs and then verify that these variables are passed into emitc.call_opaque "noc_async_write_tile" in the correct arg order instead of just using {{.*}} there? look at other tests in this file

I've given them more meaningful names as well as capturing and checking for them in later uses.

vroubtsovTT · 2025-02-27T02:38:31Z

test/ttmlir/Conversion/TTKernelToEmitC/ttkernel.mlir

+      "ttkernel.noc_async_write_tile"(%temp2, %s, %temp1) : (i32, !ttkernel.interleaved_addr_gen_fast, i32) -> ()
+      // CHECK: emitc.call_opaque "noc_async_read_tile"{{.*}}
+      "ttkernel.noc_async_read_tile"(%temp2, %s, %temp1) : (i32, !ttkernel.interleaved_addr_gen_fast, i32) -> ()
+      // CHECK-NEXT: return


this CHECK_NEXT anchor doesn't seem to be doing anything useful here

yep.. removed

vroubtsovTT · 2025-02-27T02:52:59Z

lib/Conversion/TTKernelToEmitC/TTKernelToEmitC.cpp

+      auto lvalueBankBaseAddr = rewriter.create<emitc::MemberOp>(
+          op->getLoc(),
+          emitc::LValueType::get(op.getBankBaseAddress().getType()),
+          "bank_base_address", varOp);


I think the correct pattern during dialect conversion is to obtain operand types/properties from the adaptor, not the op.

Yes, 100% agree. I used adaptor for dataformat too, not sure why I was using op here :(

vroubtsovTT · 2025-02-27T02:53:32Z

lib/Conversion/TTKernelToEmitC/TTKernelToEmitC.cpp

+          "bank_base_address", varOp);
+      auto lvaluePageSize = rewriter.create<emitc::MemberOp>(
+          op->getLoc(), emitc::LValueType::get(op.getPageSize().getType()),
+          "page_size", varOp);


same comment, shouldn't this reference adaptor, not op?

vtangTT changed the title ~~Vtang tt/pykernel new type~~ Pykernel - matmul common reader/writer kernels Feb 26, 2025

vtangTT force-pushed the vtangTT/pykernel_new_type branch from 79998c8 to 295da6a Compare February 26, 2025 21:20

vtangTT marked this pull request as ready for review February 26, 2025 21:32

vtangTT requested review from nsmithtt, tapspatel, odjuricicTT, vprajapati-tt, vcanicTT and vroubtsovTT as code owners February 26, 2025 21:32

nsmithtt approved these changes Feb 27, 2025

View reviewed changes

vroubtsovTT requested changes Feb 27, 2025

View reviewed changes

vtangTT added 17 commits February 28, 2025 03:25

Dataformat enum

2881583

InterleavedAddrGenFast draft type

0186815

wip

fb077da

addrgen kinda working, still need to support dataformat

2b50698

move the pass into TTKernel converters scope

273fde8

reader_bmm test pykernel

14b00f0

addrgen struct and get_dataformat working

c3f5a8a

noc_async_read_tile working

4517f23

reader_bmm_8bank_output_tiles_partitioned working

00254c7

reader_bmm_8bank working

4f74fdb

noc_async_write_tile working

a082190

writer_bmm_8bank working

aa24cb9

add to lit

9ee94b6

cleanup

9898d90

lit unit tests for added ops (should be working?)

d76aff4

llvm uplift fix

996b762

review

83c1e10

vtangTT force-pushed the vtangTT/pykernel_new_type branch from 78de1ce to 83c1e10 Compare February 28, 2025 03:25

vtangTT requested a review from vroubtsovTT February 28, 2025 17:47

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Pykernel - matmul common reader/writer kernels #2289

Pykernel - matmul common reader/writer kernels #2289

vtangTT commented Feb 26, 2025 •

edited

Loading

vroubtsovTT Feb 27, 2025

vtangTT Feb 28, 2025

vroubtsovTT Feb 27, 2025

vtangTT Feb 28, 2025

vroubtsovTT Feb 27, 2025

vtangTT Feb 28, 2025

vroubtsovTT Feb 27, 2025

vtangTT Feb 28, 2025

vroubtsovTT Feb 27, 2025 •

edited

Loading

vtangTT Feb 28, 2025

vroubtsovTT Feb 27, 2025

vtangTT Feb 28, 2025

vroubtsovTT Feb 27, 2025

vtangTT Feb 28, 2025

vroubtsovTT Feb 27, 2025

vtangTT Feb 28, 2025

vroubtsovTT Feb 27, 2025

vtangTT Feb 28, 2025

vroubtsovTT Feb 27, 2025

Pykernel - matmul common reader/writer kernels #2289

Are you sure you want to change the base?

Pykernel - matmul common reader/writer kernels #2289

Conversation

vtangTT commented Feb 26, 2025 • edited Loading

Ticket

Problem description

What's changed

Checklist

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

vroubtsovTT Feb 27, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

vtangTT commented Feb 26, 2025 •

edited

Loading

vroubtsovTT Feb 27, 2025 •

edited

Loading