[BUG] [CuTeDSL] Class Type resolution Error in 4.3.4/4.3.5

[HanGuo97]

Which component has the problem?

CuTe DSL

Bug Report

Describe the bug
When a class is defined inside a class and used in subsequent type annotations, latest version will raise an error saying the class is undefined. This works fine in 4.3.3 but 4.3.4, 4.3.5 will fail.

Steps/Code to reproduce bug

import cutlass
import cutlass.cute as cute

class A:
    class B:
        pass

    @cute.jit
    def f(self) -> B:
        return B()


@cute.jit
def func():
    A().f()
func()

[HanGuo97]

Perhaps loosely related, I noticed that the following code would cause the following error when I import from __future__ import annotations:

@cute.struct
class SharedStorage(EpilogueSharedStorage):
    _placeholder: cute.struct.Align[cute.struct.MemRange[cute.Float32, 0], 16]

TypeError: Struct element only support struct/array/base_dsl scalar, but got cute.struct.Align[cute.struct.MemRange[cute.Float32, 0], 16]

But the same code would run just fine without the annotations import.

I don't have a minimal example handy, but the one I have is fairly close to the dense gemm example.

[anakinxc]

Thanks for reporting, will fix.

As a workaround, remove return annotation for now

class A:
    class B:
        pass

    @cute.jit
    def f(self):
        return A.B()

[HanGuo97]

Oh, you are right... thanks!

Do you want me to keep this open?

[anakinxc]

Yes please, will fix this in next release

[HanGuo97]

ok!