Refine Python API

[hunhoffe]

This issue is a method of tracking/deconflicting efforts to change the Python API. If anyone has any PRs to change the Python API please let me know and we can record them here! We can also use this issue to record ideas/brainstorming before a PR is created.

Planned

• Clean up user error messages

In-Progress

• Use mlir-python-extras as a dependency #1828 One of my previous PRs broke some compatibility with the mlir-python-extras install. There is a workaround right now, but I still need to fix this.

Completed

• Use numpy.ndarray and numpy dtypes instead of memref types in Python API #1786
• Hide yield statements at end of python scf for-loop body #1785
• Replace npu_sync with dma_wait in programming examples #1791
• Call kernels directly in Python, index into MemRefs, other small python changes #1800
• Numpy bfloat16 workaround #1817
• Make mat mul single core python files use consistent style/API #1824
• Python buffer API refactoring, RTP Op Wrapper #1818 (includes some deduplication)
• Use ml_dtypes for bfloat16 numpy extension #1829

Notes

[fifield]

• Clean up WriteRTP/Buffer APIs e.g., code like here:

It's not clear to me that we still need the RTP ops, given that npu.write32 and npu.blockwrite offer essentially the same functionality. What's the plan for RTP ops?

[hunhoffe]

@fifield I didn't have anything major planned, I'm just trying to get rid of all the Python operations that expect symbol names instead of python objects, so I was thinking about changing

NpuWriteRTPOp("rtpComputeTile2", index=0, value=50) 

To either

rtpComputeTile2.rtpWrite(index=0, value=50)

or (not sure if this is possible to infer the right type of write operation if these is more than one, so maybe not)

rtpComputeTile2[0] = 50

Currently, NpuWriteRTPOp is used in the following programming examples:

• ml/resnet/layers_conv2_x
• ml/conv2d_fused_relu
• ml/bottleneck
• ml/conv2d
• vision/color_threshold/aie2_colorThreshold.py

I'm not sure what the differences between NpuWriteRTPOp and npu.write32/npu.blockwrite are, but alternatively I could try to port those examples to use npu.write32/npu.blockwrite instead?

[hunhoffe]

To record decisions make offline: RTP ops will probably be used with some additional synchronization method in the future, so it's worth it to clean up the Python interface