This is an archive of the discontinued LLVM Phabricator instance.

Differential D133971

[mlir][sparse][python] improve sparse encoding test
ClosedPublic

Authored by aartbik on Sep 15 2022, 1:39 PM.

Details

Reviewers
bixia
Peiming
Commits
rG0baec207ea28: [mlir][sparse][python] improve sparse encoding test

Diff Detail

Event Timeline

aartbik created this revision.Sep 15 2022, 1:39 PM
Herald added a project: Restricted Project. · View Herald TranscriptSep 15 2022, 1:39 PM
Herald added subscribers: anlunx, bzcheeseman, sdasgup3 and 18 others. · View Herald Transcript
aartbik requested review of this revision.Sep 15 2022, 1:39 PM
Herald added a project: Restricted Project. · View Herald TranscriptSep 15 2022, 1:39 PM
Herald added subscribers: stephenneuendorffer, nicolasvasilache. · View Herald Transcript
aartbik added a reviewer: Peiming.Sep 15 2022, 1:40 PM
bixia accepted this revision.Sep 15 2022, 1:55 PM
This revision is now accepted and ready to land.Sep 15 2022, 1:55 PM
Harbormaster completed remote builds in B186929: Diff 460494.Sep 15 2022, 1:56 PM
Closed by commit rG0baec207ea28: [mlir][sparse][python] improve sparse encoding test (authored by aartbik). · Explain WhySep 15 2022, 2:02 PM
This revision was automatically updated to reflect the committed changes.
aartbik added a commit: rG0baec207ea28: [mlir][sparse][python] improve sparse encoding test.

Revision Contents

PathSize
mlir/
test/
python/
dialects/
sparse_tensor/
53 lines

Diff 460505

mlir/test/python/dialects/sparse_tensor/dialect.py

# RUN: %PYTHON %s | FileCheck %s # RUN: %PYTHON %s | FileCheck %s
from mlir.ir import * from mlir.ir import *
from mlir.dialects import sparse_tensor as st from mlir.dialects import sparse_tensor as st
def run(f): def run(f):
print("\nTEST:", f.__name__) print("\nTEST:", f.__name__)
f() f()
return f return f
# CHECK-LABEL: TEST: testEncodingAttr1D # CHECK-LABEL: TEST: testEncodingAttr1D
@run @run
def testEncodingAttr1D(): def testEncodingAttr1D():
with Context() as ctx: with Context() as ctx:
parsed = Attribute.parse( parsed = Attribute.parse('#sparse_tensor.encoding<{'
'#sparse_tensor.encoding<{ dimLevelType = [ "compressed" ], ' ' dimLevelType = [ "compressed" ],'
'pointerBitWidth = 16, indexBitWidth = 32 }>') ' pointerBitWidth = 16,'
' indexBitWidth = 32'
'}>')
# CHECK: #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ], pointerBitWidth = 16, indexBitWidth = 32 }>
print(parsed) print(parsed)
casted = st.EncodingAttr(parsed) casted = st.EncodingAttr(parsed)
# CHECK: equal: True # CHECK: equal: True
print(f"equal: {casted == parsed}") print(f"equal: {casted == parsed}")
# CHECK: dim_level_types: [<DimLevelType.compressed: 1>] # CHECK: dim_level_types: [<DimLevelType.compressed: 1>]
print(f"dim_level_types: {casted.dim_level_types}") print(f"dim_level_types: {casted.dim_level_types}")
# CHECK: dim_ordering: None # CHECK: dim_ordering: None
# Note that for 1D, the ordering is None, which exercises several special
# cases.
print(f"dim_ordering: {casted.dim_ordering}") print(f"dim_ordering: {casted.dim_ordering}")
# CHECK: pointer_bit_width: 16 # CHECK: pointer_bit_width: 16
print(f"pointer_bit_width: {casted.pointer_bit_width}") print(f"pointer_bit_width: {casted.pointer_bit_width}")
# CHECK: index_bit_width: 32 # CHECK: index_bit_width: 32
print(f"index_bit_width: {casted.index_bit_width}") print(f"index_bit_width: {casted.index_bit_width}")
created = st.EncodingAttr.get(casted.dim_level_types, None, 16, 32) created = st.EncodingAttr.get(casted.dim_level_types, None, 0, 0)
# CHECK: #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
print(created) print(created)
# CHECK: created_equal: True # CHECK: created_equal: False
print(f"created_equal: {created == casted}") print(f"created_equal: {created == casted}")
# Verify that the factory creates an instance of the proper type. # Verify that the factory creates an instance of the proper type.
# CHECK: is_proper_instance: True # CHECK: is_proper_instance: True
print(f"is_proper_instance: {isinstance(created, st.EncodingAttr)}") print(f"is_proper_instance: {isinstance(created, st.EncodingAttr)}")
# CHECK: created_pointer_bit_width: 16 # CHECK: created_pointer_bit_width: 0
print(f"created_pointer_bit_width: {created.pointer_bit_width}") print(f"created_pointer_bit_width: {created.pointer_bit_width}")
# CHECK-LABEL: TEST: testEncodingAttr2D # CHECK-LABEL: TEST: testEncodingAttr2D
@run @run
def testEncodingAttr2D(): def testEncodingAttr2D():
with Context() as ctx: with Context() as ctx:
parsed = Attribute.parse( parsed = Attribute.parse('#sparse_tensor.encoding<{'
'#sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], ' ' dimLevelType = [ "dense", "compressed" ],'
'dimOrdering = affine_map<(d0, d1) -> (d0, d1)>, ' ' dimOrdering = affine_map<(d0, d1) -> (d1, d0)>,'
'pointerBitWidth = 16, indexBitWidth = 32 }>') ' pointerBitWidth = 8,'
' indexBitWidth = 32'
'}>')
# CHECK: #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)>, pointerBitWidth = 8, indexBitWidth = 32 }>
print(parsed) print(parsed)
casted = st.EncodingAttr(parsed) casted = st.EncodingAttr(parsed)
# CHECK: equal: True # CHECK: equal: True
print(f"equal: {casted == parsed}") print(f"equal: {casted == parsed}")
# CHECK: dim_level_types: [<DimLevelType.dense: 0>, <DimLevelType.compressed: 1>] # CHECK: dim_level_types: [<DimLevelType.dense: 0>, <DimLevelType.compressed: 1>]
print(f"dim_level_types: {casted.dim_level_types}") print(f"dim_level_types: {casted.dim_level_types}")
# CHECK: dim_ordering: (d0, d1) -> (d0, d1) # CHECK: dim_ordering: (d0, d1) -> (d1, d0)
print(f"dim_ordering: {casted.dim_ordering}") print(f"dim_ordering: {casted.dim_ordering}")
# CHECK: pointer_bit_width: 16 # CHECK: pointer_bit_width: 8
print(f"pointer_bit_width: {casted.pointer_bit_width}") print(f"pointer_bit_width: {casted.pointer_bit_width}")
# CHECK: index_bit_width: 32 # CHECK: index_bit_width: 32
print(f"index_bit_width: {casted.index_bit_width}") print(f"index_bit_width: {casted.index_bit_width}")
created = st.EncodingAttr.get(casted.dim_level_types, casted.dim_ordering, created = st.EncodingAttr.get(casted.dim_level_types, casted.dim_ordering,
16, 32) 8, 32)
# CHECK: #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)>, pointerBitWidth = 8, indexBitWidth = 32 }>
print(created) print(created)
# CHECK: created_equal: True # CHECK: created_equal: True
print(f"created_equal: {created == casted}") print(f"created_equal: {created == casted}")
# CHECK-LABEL: TEST: testEncodingAttrOnTensor # CHECK-LABEL: TEST: testEncodingAttrOnTensorType
@run @run
def testEncodingAttrOnTensor(): def testEncodingAttrOnTensorType():
with Context() as ctx, Location.unknown(): with Context() as ctx, Location.unknown():
encoding = st.EncodingAttr(Attribute.parse( encoding = st.EncodingAttr(
'#sparse_tensor.encoding<{ dimLevelType = [ "compressed" ], ' Attribute.parse('#sparse_tensor.encoding<{'
'pointerBitWidth = 16, indexBitWidth = 32 }>')) ' dimLevelType = [ "compressed" ], '
' pointerBitWidth = 64,'
' indexBitWidth = 32'
'}>'))
tt = RankedTensorType.get((1024,), F32Type.get(), encoding=encoding) tt = RankedTensorType.get((1024,), F32Type.get(), encoding=encoding)
# CHECK: tensor<1024xf32, #sparse_tensor # CHECK: tensor<1024xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ], pointerBitWidth = 64, indexBitWidth = 32 }>>
print(tt) print(tt)
# CHECK: #sparse_tensor.encoding # CHECK: #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ], pointerBitWidth = 64, indexBitWidth = 32 }>
print(tt.encoding) print(tt.encoding)
assert tt.encoding == encoding assert tt.encoding == encoding