[Bug] TMulSOp，When the first param is not F32, the second value follow the first param, the python gives out error

[yaomicat]

Component

Python bindings (python/)

Description

【issue】TMulSOp the second param must be F32, which is not expected

Reproduction (minimal)

from mlir.ir import Context, Location, Module, InsertionPoint
from mlir.dialects import func, arith, scf, pto
from mlir.ir import IndexType
from mlir.ir import F16Type


def build():
    with Context() as ctx:
        pto.register_dialect(ctx, load=True)

        with Location.unknown(ctx):
            m = Module.create()

            idx = IndexType.get(ctx)
            f16 = F16Type.get(ctx)
            ptr_f16 = pto.PtrType.get(f16, ctx)
            tv2_f16 = pto.TensorViewType.get(2, f16, ctx)
            tile_view_1x16_f16 = pto.PartitionTensorViewType.get([1, 16], f16, ctx)
            vec = pto.AddressSpaceAttr.get(pto.AddressSpace.VEC, ctx)
            bl = pto.BLayoutAttr.get(pto.BLayout.RowMajor, ctx)
            sl = pto.SLayoutAttr.get(pto.SLayout.NoneBox, ctx)
            pd = pto.PadValueAttr.get(pto.PadValue.Null, ctx)
            fractal_ab_size = pto.TileConfig.fractalABSize
            cfg = pto.TileBufConfigAttr.get(bl, sl, fractal_ab_size, pd, ctx)
            tile_buf_1x16_f16 = pto.TileBufType.get([1, 16], f16, vec, [1, 16], cfg, ctx)

            fn_ty = func.FunctionType.get([ptr_f16, ptr_f16], [])
            with InsertionPoint(m.body):
                fn = func.FuncOp("generated_kernel", fn_ty)
                entry = fn.add_entry_block()

            with InsertionPoint(entry):
                c0 = arith.ConstantOp(idx, 0).result
                c1 = arith.ConstantOp(idx, 1).result
                c2 = arith.ConstantOp(idx, 2).result
                c16 = arith.ConstantOp(idx, 16).result
                args = list(entry.arguments)
                src0 = args[0]
                dst = args[1]

                tv0 = pto.MakeTensorViewOp(tv2_f16, src0, [c1, c16], [c16, c1]).result
                pv0 = pto.PartitionViewOp(tile_view_1x16_f16, tv0, offsets=[c0, c0], sizes=[c1, c16]).result
                tv1 = pto.MakeTensorViewOp(tv2_f16, dst, [c1, c16], [c16, c1]).result
                pv1 = pto.PartitionViewOp(tile_view_1x16_f16, tv1, offsets=[c0, c0], sizes=[c1, c16]).result

                loop = scf.ForOp(c0, c2, c1, [])
                with InsertionPoint(loop.body):
                    tb0 = pto.AllocTileOp(tile_buf_1x16_f16).result
                    tb1 = pto.AllocTileOp(tile_buf_1x16_f16).result

                    pto.TLoadOp(None, pv0, tb0)

                    scale = arith.ConstantOp(f16, 1.0).result

                    pto.TMulSOp(tb0, scale, tb1)

                    pto.TStoreOp(None, tb1, pv1)
                    scf.YieldOp([])

                func.ReturnOp([])

            m.operation.verify()
            return m


if __name__ == "__main__":
    print(build())

Expected behavior

compile pass

python test.py

Actual behavior / error logs

$ python TMulSOp_f16_1_16.py
Traceback (most recent call last):
  File "/home/cplop/code/pto_test/tools/ptoastest/test/TMulSOp/TMulSOp_f16_1_16.py", line 67, in <module>
    print(build())
          ^^^^^^^
  File "/home/cplop/code/pto_test/tools/ptoastest/test/TMulSOp/TMulSOp_f16_1_16.py", line 62, in build
    m.operation.verify()
mlir._mlir_libs._site_initialize.<locals>.MLIRError: Verification failed:
error: unknown: 'pto.tmuls' op operand #1 must be 32-bit float, but got 'f16'
 note: unknown: see current operation: "pto.tmuls"(%8, %10, %9) : (!pto.tile_buf<loc=vec, dtype=f16, rows=1, cols=16, v_row=1, v_col=16, blayout=row_major, slayout=none_box, fractal=512, pad=0>, f16, !pto.tile_buf<loc=vec, dtype=f16, rows=1, cols=16, v_row=1, v_col=16, blayout=row_major, slayout=none_box, fractal=512, pad=0>) -> ()

Git commit

bba77f1

Host platform

None

Target Ascend arch (if relevant)

None

PTOAS build level (if relevant)

None