[Issue]: make_layout_tv API fails due to compilation issue

[rahulbatra85]

Problem Description

fx.make_layout_tv API fails due to compilation issue. I see unit tests are missing for this API.

Operating System

Ubuntu 24.04.3 LTS (Noble Numbat)

CPU

AMD EPYC 9575F 64-Core Processor

GPU

AMD Instinct MI355, gfx950

ROCm Version

ROCm 7.1.0

ROCm Component

No response

Steps to Reproduce

import torch

import flydsl.compiler as flyc
import flydsl.expr as fx

@flyc.jit
def tv_layout(
    mA: fx.Tensor,
):
    thr_layout = fx.make_layout((4,64), (64,1))
    val_layout = fx.make_layout((4,8), (8,1))
    tiler_mn, tv_layout = fx.make_layout_tv(thr_layout, val_layout)

M, N = (2048, 2048)
A = torch.randn(M,N,dtype=torch.bfloat16).cuda()
tv_layout(A)

python3 examples/tv_layout.py 
Traceback (most recent call last):
  File "/workspace/FlyDSL2/examples/tv_layout.py", line 16, in <module>
    tv_layout(A)
  File "/workspace/FlyDSL2/python/flydsl/compiler/jit_function.py", line 816, in __call__
    compiled_module = MlirCompiler.compile(module, arch=backend.target.arch, func_name=self.func.__name__)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/FlyDSL2/python/flydsl/compiler/jit_function.py", line 372, in compile
    module.operation.verify()
flydsl._mlir._mlir_libs._mlir.ir.MLIRError: Verification failed:
error: "tmp = make_layout((thr_size, val_size), (1, thr_size))"("/workspace/FlyDSL2/python/flydsl/expr/derived.py":169:10): 'fly.make_shape' op operand #0 must be variadic of 32-bit signless integer or 64-bit signless integer, but got '!fly.int_tuple<256>'
 note: "tmp = make_layout((thr_size, val_size), (1, thr_size))"("/workspace/FlyDSL2/python/flydsl/expr/derived.py":169:10): see current operation: %9 = "fly.make_shape"(%7, %8) : (!fly.int_tuple<256>, !fly.int_tuple<32>) -> !fly.int_tuple<(?,?)>

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

No response

Additional Information

No response