Something went wrong. Try again.
What happens if I push all of nixpkgs
Something went wrong. Try again.
241 kB · 4617 lines
at litex
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788278927902791279227932794279527962797279827992800280128022803280428052806280728082809281028112812281328142815281628172818281928202821282228232824282528262827282828292830283128322833283428352836283728382839284028412842284328442845284628472848284928502851285228532854285528562857285828592860286128622863286428652866286728682869287028712872287328742875287628772878287928802881288228832884288528862887288828892890289128922893289428952896289728982899290029012902290329042905290629072908290929102911291229132914291529162917291829192920292129222923292429252926292729282929293029312932293329342935293629372938293929402941294229432944294529462947294829492950295129522953295429552956295729582959296029612962296329642965296629672968296929702971297229732974297529762977297829792980298129822983298429852986298729882989299029912992299329942995299629972998299930003001300230033004300530063007300830093010301130123013301430153016301730183019302030213022302330243025302630273028302930303031303230333034303530363037303830393040304130423043304430453046304730483049305030513052305330543055305630573058305930603061306230633064306530663067306830693070307130723073307430753076307730783079308030813082308330843085308630873088308930903091309230933094309530963097309830993100310131023103310431053106310731083109311031113112311331143115311631173118311931203121312231233124312531263127312831293130313131323133313431353136313731383139314031413142314331443145314631473148314931503151315231533154315531563157315831593160316131623163316431653166316731683169317031713172317331743175317631773178317931803181318231833184318531863187318831893190319131923193319431953196319731983199320032013202320332043205320632073208320932103211321232133214321532163217321832193220322132223223322432253226322732283229323032313232323332343235323632373238323932403241324232433244324532463247324832493250325132523253325432553256325732583259326032613262326332643265326632673268326932703271327232733274327532763277327832793280328132823283328432853286328732883289329032913292329332943295329632973298329933003301330233033304330533063307330833093310331133123313331433153316331733183319332033213322332333243325332633273328332933303331333233333334333533363337333833393340334133423343334433453346334733483349335033513352335333543355335633573358335933603361336233633364336533663367336833693370337133723373337433753376337733783379338033813382338333843385338633873388338933903391339233933394339533963397339833993400340134023403340434053406340734083409341034113412341334143415341634173418341934203421342234233424342534263427342834293430343134323433343434353436343734383439344034413442344334443445344634473448344934503451345234533454345534563457345834593460346134623463346434653466346734683469347034713472347334743475347634773478347934803481348234833484348534863487348834893490349134923493349434953496349734983499350035013502350335043505350635073508350935103511351235133514351535163517351835193520352135223523352435253526352735283529353035313532353335343535353635373538353935403541354235433544354535463547354835493550355135523553355435553556355735583559356035613562356335643565356635673568356935703571357235733574357535763577357835793580358135823583358435853586358735883589359035913592359335943595359635973598359936003601360236033604360536063607360836093610361136123613361436153616361736183619362036213622362336243625362636273628362936303631363236333634363536363637363836393640364136423643364436453646364736483649365036513652365336543655365636573658365936603661366236633664366536663667366836693670367136723673367436753676367736783679368036813682368336843685368636873688368936903691369236933694369536963697369836993700370137023703370437053706370737083709371037113712371337143715371637173718371937203721372237233724372537263727372837293730373137323733373437353736373737383739374037413742374337443745374637473748374937503751375237533754375537563757375837593760376137623763376437653766376737683769377037713772377337743775377637773778377937803781378237833784378537863787378837893790379137923793379437953796379737983799380038013802380338043805380638073808380938103811381238133814381538163817381838193820382138223823382438253826382738283829383038313832383338343835383638373838383938403841384238433844384538463847384838493850385138523853385438553856385738583859386038613862386338643865386638673868386938703871387238733874387538763877387838793880388138823883388438853886388738883889389038913892389338943895389638973898389939003901390239033904390539063907390839093910391139123913391439153916391739183919392039213922392339243925392639273928392939303931393239333934393539363937393839393940394139423943394439453946394739483949395039513952395339543955395639573958395939603961396239633964396539663967396839693970397139723973397439753976397739783979398039813982398339843985398639873988398939903991399239933994399539963997399839994000400140024003400440054006400740084009401040114012401340144015401640174018401940204021402240234024402540264027402840294030403140324033403440354036403740384039404040414042404340444045404640474048404940504051405240534054405540564057405840594060406140624063406440654066406740684069407040714072407340744075407640774078407940804081408240834084408540864087408840894090409140924093409440954096409740984099410041014102410341044105410641074108410941104111411241134114411541164117411841194120412141224123412441254126412741284129413041314132413341344135413641374138413941404141414241434144414541464147414841494150415141524153415441554156415741584159416041614162416341644165416641674168416941704171417241734174417541764177417841794180418141824183418441854186418741884189419041914192419341944195419641974198419942004201420242034204420542064207420842094210421142124213421442154216421742184219422042214222422342244225422642274228422942304231423242334234423542364237423842394240424142424243424442454246424742484249425042514252425342544255425642574258425942604261426242634264426542664267426842694270427142724273427442754276427742784279428042814282428342844285428642874288428942904291429242934294429542964297429842994300430143024303430443054306430743084309431043114312431343144315431643174318431943204321432243234324432543264327432843294330433143324333433443354336433743384339434043414342434343444345434643474348434943504351435243534354435543564357435843594360436143624363436443654366436743684369437043714372437343744375437643774378437943804381438243834384438543864387438843894390439143924393439443954396439743984399440044014402440344044405440644074408440944104411441244134414441544164417441844194420442144224423442444254426442744284429443044314432443344344435443644374438443944404441444244434444444544464447444844494450445144524453445444554456445744584459446044614462446344644465446644674468446944704471447244734474447544764477447844794480448144824483448444854486448744884489449044914492449344944495449644974498449945004501450245034504450545064507450845094510451145124513451445154516451745184519452045214522452345244525452645274528452945304531453245334534453545364537453845394540454145424543454445454546454745484549455045514552455345544555455645574558455945604561456245634564456545664567456845694570457145724573457445754576457745784579458045814582458345844585458645874588458945904591459245934594459545964597459845994600460146024603460446054606460746084609461046114612461346144615461646174618From fc7a8e35819bda632bdcf1cf75fd9abe4d4e077a Mon Sep 17 00:00:00 2001From: Christian Sigg <chsigg@users.noreply.github.com>Date: Thu, 16 Feb 2023 15:40:53 +0100Subject: [PATCH] Rebase Triton to LLVM-15. (#1070)
This PR rebases Triton from LLVM-14 to LLVM-15. Most changes aremechanical, except for the analysis framework changes.--- CMakeLists.txt | 6 +- bin/CMakeLists.txt | 2 +- bin/FileCheck/FileCheck.cpp | 3 + bin/triton-opt.cpp | 6 +- bin/triton-translate.cpp | 7 +- include/triton/Analysis/Alias.h | 21 +- include/triton/Analysis/Allocation.h | 2 + include/triton/Analysis/AxisInfo.h | 56 ++- include/triton/Analysis/Utility.h | 6 +- include/triton/Conversion/Passes.td | 4 +- include/triton/Dialect/Triton/IR/Dialect.h | 7 +- .../triton/Dialect/Triton/IR/TritonDialect.td | 8 +- include/triton/Dialect/Triton/IR/TritonOps.td | 12 +- .../triton/Dialect/Triton/IR/TritonTypes.td | 2 + .../Dialect/Triton/Transforms/Passes.td | 3 +- include/triton/Dialect/TritonGPU/IR/Dialect.h | 4 +- .../Dialect/TritonGPU/IR/TritonGPUAttrDefs.td | 7 + .../Dialect/TritonGPU/IR/TritonGPUDialect.td | 2 +- .../Dialect/TritonGPU/IR/TritonGPUOps.td | 13 +- lib/Analysis/Alias.cpp | 14 +- lib/Analysis/Allocation.cpp | 30 +- lib/Analysis/AxisInfo.cpp | 79 ++-- lib/Analysis/CMakeLists.txt | 2 +- lib/Analysis/Membar.cpp | 2 +- lib/Analysis/Utility.cpp | 54 +++ .../TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp | 3 - lib/Conversion/TritonGPUToLLVM/DotOpHelpers.h | 10 +- .../TritonGPUToLLVM/DotOpToLLVM.cpp | 5 - .../TritonGPUToLLVM/ElementwiseOpToLLVM.cpp | 2 - .../TritonGPUToLLVM/LoadStoreOpToLLVM.cpp | 5 +- .../TritonGPUToLLVM/ReduceOpToLLVM.cpp | 2 - .../TritonGPUToLLVM/TritonGPUToLLVM.cpp | 7 +- .../TritonGPUToLLVM/TritonGPUToLLVMBase.h | 26 +- .../TritonGPUToLLVM/TritonGPUToLLVMPass.cpp | 52 +-- lib/Conversion/TritonGPUToLLVM/Utility.h | 5 +- .../TritonToTritonGPUPass.cpp | 69 ++-- lib/Dialect/Triton/IR/CMakeLists.txt | 10 +- lib/Dialect/Triton/IR/Ops.cpp | 34 +- lib/Dialect/Triton/Transforms/Combine.cpp | 6 +- lib/Dialect/Triton/Transforms/Combine.td | 2 +- lib/Dialect/TritonGPU/IR/Dialect.cpp | 27 +- lib/Dialect/TritonGPU/Transforms/Coalesce.cpp | 20 +- lib/Dialect/TritonGPU/Transforms/Combine.cpp | 2 +- lib/Dialect/TritonGPU/Transforms/Combine.td | 1 + .../Transforms/DecomposeConversions.cpp | 2 +- lib/Dialect/TritonGPU/Transforms/Pipeline.cpp | 10 +- .../Transforms/ReorderInstructions.cpp | 2 +- .../Transforms/TritonGPUConversion.cpp | 12 +- .../Transforms/UpdateMmaForVolta.cpp | 6 +- lib/Dialect/TritonGPU/Transforms/Utility.cpp | 2 +- lib/Target/LLVMIR/CMakeLists.txt | 3 +- lib/Target/PTX/PTXTranslation.cpp | 3 + python/setup.py | 15 +- python/src/triton.cc | 85 +++-- python/test/unit/language/test_core.py | 2 +- python/triton/compiler.py | 4 +- test/Analysis/test-alias.mlir | 24 +- test/Analysis/test-alignment.mlir | 344 +++++++++--------- test/Analysis/test-allocation.mlir | 32 +- test/Analysis/test-membar.mlir | 38 +- test/Conversion/triton_ops.mlir | 10 +- test/Conversion/triton_to_tritongpu.mlir | 6 +- test/Conversion/tritongpu_to_llvm.mlir | 94 ++--- test/Target/tritongpu_to_llvmir.mlir | 4 +- test/Target/tritongpu_to_ptx.mlir | 2 +- test/Triton/combine.mlir | 40 +- test/Triton/vecadd.mlir | 4 +- test/TritonGPU/coalesce.mlir | 2 +- test/TritonGPU/combine.mlir | 38 +- test/TritonGPU/loop-pipeline.mlir | 22 +- test/TritonGPU/matmul.mlir | 4 +- test/TritonGPU/prefetch.mlir | 4 +- test/TritonGPU/update-mma-for-volta.mlir | 4 +- test/lib/Analysis/TestAlias.cpp | 29 +- test/lib/Analysis/TestAllocation.cpp | 5 +- test/lib/Analysis/TestAxisInfo.cpp | 51 +-- test/lib/Analysis/TestMembar.cpp | 7 +- 78 files changed, 808 insertions(+), 742 deletions(-)
diff --git a/CMakeLists.txt b/CMakeLists.txtindex d0d361fc7c..b281a28400 100644--- a/CMakeLists.txt+++ b/CMakeLists.txt@@ -1,4 +1,7 @@ cmake_minimum_required(VERSION 3.6)++cmake_policy(SET CMP0116 OLD)+ include(ExternalProject) set(CMAKE_CXX_STANDARD 17)@@ -155,7 +158,6 @@ if(TRITON_BUILD_PYTHON_MODULE) endif() endif() - # # Triton # file(GLOB_RECURSE LIBTRITON_SRC lib/*.cc) # if (WIN32 AND TRITON_BUILD_PYTHON_MODULE)@@ -212,7 +214,7 @@ if(TRITON_BUILD_PYTHON_MODULE) # optimizations MLIRPass MLIRTransforms- MLIRLLVMIR+ MLIRLLVMDialect MLIRSupport MLIRTargetLLVMIRExport MLIRExecutionEnginediff --git a/bin/CMakeLists.txt b/bin/CMakeLists.txtindex 906f635f8b..695b3479fd 100644--- a/bin/CMakeLists.txt+++ b/bin/CMakeLists.txt@@ -48,7 +48,7 @@ llvm_update_compile_flags(triton-translate) # MLIR core MLIROptLib MLIRIR- MLIRLLVMIR+ MLIRLLVMDialect MLIRPass MLIRSupport MLIRTransformsdiff --git a/bin/FileCheck/FileCheck.cpp b/bin/FileCheck/FileCheck.cppindex 819efc3541..9ac6f1b277 100644--- a/bin/FileCheck/FileCheck.cpp+++ b/bin/FileCheck/FileCheck.cpp@@ -19,6 +19,7 @@ #include "llvm/Support/CommandLine.h" #include "llvm/Support/InitLLVM.h" #include "llvm/Support/Process.h"+#include "llvm/Support/SourceMgr.h" #include "llvm/Support/WithColor.h" #include "llvm/Support/raw_ostream.h" #include <cmath>@@ -360,6 +361,8 @@ static std::string GetCheckTypeAbbreviation(Check::FileCheckType Ty) { return "bad-not"; case Check::CheckBadCount: return "bad-count";+ case Check::CheckMisspelled:+ return "misspelled"; case Check::CheckNone: llvm_unreachable("invalid FileCheckType"); }diff --git a/bin/triton-opt.cpp b/bin/triton-opt.cppindex 9f3b53b7ae..f96232e1b0 100644--- a/bin/triton-opt.cpp+++ b/bin/triton-opt.cpp@@ -8,7 +8,7 @@ #include "mlir/IR/Dialect.h" #include "mlir/InitAllPasses.h"-#include "mlir/Support/MlirOptMain.h"+#include "mlir/Tools/mlir-opt/MlirOptMain.h" namespace mlir { namespace test {@@ -33,8 +33,8 @@ int main(int argc, char **argv) { // TODO: register Triton & TritonGPU passes mlir::DialectRegistry registry; registry.insert<mlir::triton::TritonDialect,- mlir::triton::gpu::TritonGPUDialect, mlir::math::MathDialect,- mlir::arith::ArithmeticDialect, mlir::StandardOpsDialect,+ mlir::triton::gpu::TritonGPUDialect, mlir::func::FuncDialect,+ mlir::math::MathDialect, mlir::arith::ArithmeticDialect, mlir::scf::SCFDialect, mlir::gpu::GPUDialect>(); return mlir::asMainReturnCode(mlir::MlirOptMain(diff --git a/bin/triton-translate.cpp b/bin/triton-translate.cppindex 05ba15e453..56b5d65857 100644--- a/bin/triton-translate.cpp+++ b/bin/triton-translate.cpp@@ -3,7 +3,7 @@ #include "mlir/IR/AsmState.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/Dialect.h"-#include "mlir/Parser.h"+#include "mlir/Parser/Parser.h" #include "mlir/Pass/Pass.h" #include "mlir/Pass/PassManager.h" #include "mlir/Support/FileUtilities.h"@@ -38,7 +38,7 @@ OwningOpRef<ModuleOp> loadMLIRModule(llvm::StringRef inputFilename, mlir::DialectRegistry registry; registry.insert<TritonDialect, triton::gpu::TritonGPUDialect, mlir::math::MathDialect, arith::ArithmeticDialect,- StandardOpsDialect, scf::SCFDialect>();+ scf::SCFDialect>(); context.appendDialectRegistry(registry); @@ -50,7 +50,8 @@ OwningOpRef<ModuleOp> loadMLIRModule(llvm::StringRef inputFilename, context.loadAllAvailableDialects(); context.allowUnregisteredDialects(); - OwningOpRef<ModuleOp> module(parseSourceFile(sourceMgr, &context));+ OwningOpRef<ModuleOp> module =+ parseSourceFile<ModuleOp>(sourceMgr, &context); if (!module) { llvm::errs() << "Parse MLIR file failed."; return nullptr;diff --git a/include/triton/Analysis/Alias.h b/include/triton/Analysis/Alias.hindex fa6b906fc9..631df518bc 100644--- a/include/triton/Analysis/Alias.h+++ b/include/triton/Analysis/Alias.h@@ -2,7 +2,7 @@ #define TRITON_ANALYSIS_ALIAS_H #include "mlir/Analysis/AliasAnalysis.h"-#include "mlir/Analysis/DataFlowAnalysis.h"+#include "mlir/Analysis/DataFlow/SparseAnalysis.h" #include "llvm/ADT/DenseSet.h" namespace mlir {@@ -21,7 +21,7 @@ class AliasInfo { } /// The pessimistic value state of a value without alias- static AliasInfo getPessimisticValueState(MLIRContext *context) {+ static AliasInfo getPessimisticValueState(MLIRContext *context = nullptr) { return AliasInfo(); } static AliasInfo getPessimisticValueState(Value value) { return AliasInfo(); }@@ -29,6 +29,10 @@ class AliasInfo { /// The union of both arguments static AliasInfo join(const AliasInfo &lhs, const AliasInfo &rhs); + void print(raw_ostream &os) const {+ llvm::interleaveComma(allocs, os, [&](Value alloc) { alloc.print(os); });+ }+ private: /// The set of allocated values that are aliased by this lattice. /// For now, we only consider aliased value produced by the following@@ -58,9 +62,13 @@ class AliasInfo { //===----------------------------------------------------------------------===// // Shared Memory Alias Analysis //===----------------------------------------------------------------------===//-class SharedMemoryAliasAnalysis : public ForwardDataFlowAnalysis<AliasInfo> {+class SharedMemoryAliasAnalysis+ : public dataflow::SparseDataFlowAnalysis<dataflow::Lattice<AliasInfo>> { public:- using ForwardDataFlowAnalysis<AliasInfo>::ForwardDataFlowAnalysis;+ using dataflow::SparseDataFlowAnalysis<+ dataflow::Lattice<AliasInfo>>::SparseDataFlowAnalysis;+ using dataflow::SparseDataFlowAnalysis<+ dataflow::Lattice<AliasInfo>>::getLatticeElement; /// XXX(Keren): Compatible interface with MLIR AliasAnalysis for future use. /// Given two values, returns their aliasing behavior.@@ -70,9 +78,10 @@ class SharedMemoryAliasAnalysis : public ForwardDataFlowAnalysis<AliasInfo> { ModRefResult getModRef(Operation *op, Value location); /// Computes if the alloc set of the results are changed.- ChangeResult+ void visitOperation(Operation *op,- ArrayRef<LatticeElement<AliasInfo> *> operands) override;+ ArrayRef<const dataflow::Lattice<AliasInfo> *> operands,+ ArrayRef<dataflow::Lattice<AliasInfo> *> results) override; }; } // namespace mlirdiff --git a/include/triton/Analysis/Allocation.h b/include/triton/Analysis/Allocation.hindex b7c136d602..89b77034cc 100644--- a/include/triton/Analysis/Allocation.h+++ b/include/triton/Analysis/Allocation.h@@ -188,6 +188,8 @@ class Allocation { friend class triton::AllocationAnalysis; }; +template <typename T> Interval(T, T) -> Interval<T>;+ } // namespace mlir #endif // TRITON_ANALYSIS_ALLOCATION_Hdiff --git a/include/triton/Analysis/AxisInfo.h b/include/triton/Analysis/AxisInfo.hindex fdfbd8fbb3..7083b9c43b 100644--- a/include/triton/Analysis/AxisInfo.h+++ b/include/triton/Analysis/AxisInfo.h@@ -1,9 +1,10 @@ #ifndef TRITON_ANALYSIS_AXISINFO_H #define TRITON_ANALYSIS_AXISINFO_H -#include "mlir/Analysis/DataFlowAnalysis.h"+#include "mlir/Analysis/DataFlow/SparseAnalysis.h" #include "llvm/Support/raw_ostream.h" +#include "mlir/Support/LLVM.h" #include "triton/Analysis/Utility.h" #include "triton/Dialect/Triton/IR/Dialect.h" #include "triton/Dialect/TritonGPU/IR/Dialect.h"@@ -62,7 +63,7 @@ class AxisInfo { } /// The pessimistic value state of the contiguity is unknown.- static AxisInfo getPessimisticValueState(MLIRContext *context) {+ static AxisInfo getPessimisticValueState(MLIRContext *context = nullptr) { return AxisInfo(); } static AxisInfo getPessimisticValueState(Value value);@@ -70,6 +71,22 @@ class AxisInfo { /// The gcd of both arguments for each dimension static AxisInfo join(const AxisInfo &lhs, const AxisInfo &rhs); + void print(raw_ostream &os) const {+ auto print = [&](StringRef name, DimVectorT vec) {+ os << name << " = [";+ llvm::interleaveComma(vec, os);+ os << "]";+ };+ print("contiguity", contiguity);+ print(", divisibility", divisibility);+ print(", constancy", constancy);+ os << ", constant_value = ";+ if (constantValue)+ os << *constantValue;+ else+ os << "<none>";+ }+ private: /// The _contiguity_ information maps the `d`-th /// dimension to the length of the shortest@@ -147,7 +164,8 @@ class AxisInfoVisitor { } virtual AxisInfo- getAxisInfo(Operation *op, ArrayRef<LatticeElement<AxisInfo> *> operands) = 0;+ getAxisInfo(Operation *op,+ ArrayRef<const dataflow::Lattice<AxisInfo> *> operands) = 0; virtual bool match(Operation *op) = 0; };@@ -157,15 +175,16 @@ template <typename OpTy> class AxisInfoVisitorImpl : public AxisInfoVisitor { public: using AxisInfoVisitor::AxisInfoVisitor; - AxisInfo getAxisInfo(Operation *op,- ArrayRef<LatticeElement<AxisInfo> *> operands) final {+ AxisInfo+ getAxisInfo(Operation *op,+ ArrayRef<const dataflow::Lattice<AxisInfo> *> operands) final { return getAxisInfo(cast<OpTy>(op), operands); } bool match(Operation *op) final { return isa<OpTy>(op); } - virtual AxisInfo getAxisInfo(OpTy op,- ArrayRef<LatticeElement<AxisInfo> *> operands) {+ virtual AxisInfo+ getAxisInfo(OpTy op, ArrayRef<const dataflow::Lattice<AxisInfo> *> operands) { llvm_unreachable("Unimplemented getAxisInfo"); } };@@ -176,8 +195,9 @@ class BinaryOpVisitorImpl : public AxisInfoVisitorImpl<OpTy> { public: using AxisInfoVisitorImpl<OpTy>::AxisInfoVisitorImpl; - AxisInfo getAxisInfo(OpTy op,- ArrayRef<LatticeElement<AxisInfo> *> operands) override {+ AxisInfo+ getAxisInfo(OpTy op,+ ArrayRef<const dataflow::Lattice<AxisInfo> *> operands) override { auto lhsInfo = operands[0]->getValue(); auto rhsInfo = operands[1]->getValue(); auto rank = lhsInfo.getRank();@@ -230,7 +250,8 @@ class AxisInfoVisitorList { (visitors.emplace_back(std::make_unique<Ts>()), ...); } - AxisInfo apply(Operation *op, ArrayRef<LatticeElement<AxisInfo> *> operands) {+ AxisInfo apply(Operation *op,+ ArrayRef<const dataflow::Lattice<AxisInfo> *> operands) { for (auto &visitor : visitors) if (visitor->match(op)) return visitor->getAxisInfo(op, operands);@@ -241,16 +262,19 @@ class AxisInfoVisitorList { std::vector<std::unique_ptr<AxisInfoVisitor>> visitors; }; -class AxisInfoAnalysis : public ForwardDataFlowAnalysis<AxisInfo> {+class AxisInfoAnalysis+ : public dataflow::SparseDataFlowAnalysis<dataflow::Lattice<AxisInfo>> { private: AxisInfoVisitorList visitors; public:- AxisInfoAnalysis(MLIRContext *context);+ AxisInfoAnalysis(DataFlowSolver &solver);+ using dataflow::SparseDataFlowAnalysis<+ dataflow::Lattice<AxisInfo>>::getLatticeElement; - ChangeResult- visitOperation(Operation *op,- ArrayRef<LatticeElement<AxisInfo> *> operands) override;+ void visitOperation(Operation *op,+ ArrayRef<const dataflow::Lattice<AxisInfo> *> operands,+ ArrayRef<dataflow::Lattice<AxisInfo> *> results) override; unsigned getPtrContiguity(Value ptr); @@ -261,4 +285,4 @@ class AxisInfoAnalysis : public ForwardDataFlowAnalysis<AxisInfo> { } // namespace mlir -#endif\ No newline at end of file+#endifdiff --git a/include/triton/Analysis/Utility.h b/include/triton/Analysis/Utility.hindex c5ac137dc1..ee7fadb59d 100644--- a/include/triton/Analysis/Utility.h+++ b/include/triton/Analysis/Utility.h@@ -1,6 +1,7 @@ #ifndef TRITON_ANALYSIS_UTILITY_H #define TRITON_ANALYSIS_UTILITY_H +#include "mlir/Analysis/DataFlowFramework.h" #include "mlir/Analysis/SliceAnalysis.h" #include "triton/Dialect/TritonGPU/IR/Dialect.h" #include <algorithm>@@ -12,7 +13,7 @@ namespace mlir { class ReduceOpHelper { public: explicit ReduceOpHelper(triton::ReduceOp op) : op(op) {- srcTy = op.operand().getType().cast<RankedTensorType>();+ srcTy = op.getOperand().getType().cast<RankedTensorType>(); } ArrayRef<int64_t> getSrcShape() { return srcTy.getShape(); }@@ -103,6 +104,9 @@ SetVector<Operation *> multiRootGetSlice(Operation *op, TransitiveFilter backwardFilter = nullptr, TransitiveFilter forwardFilter = nullptr); +// Create a basic DataFlowSolver with constant and dead code analysis included.+std::unique_ptr<DataFlowSolver> createDataFlowSolver();+ } // namespace mlir #endif // TRITON_ANALYSIS_UTILITY_Hdiff --git a/include/triton/Conversion/Passes.td b/include/triton/Conversion/Passes.tdindex 70bb20b78e..be00eb2dac 100644--- a/include/triton/Conversion/Passes.td+++ b/include/triton/Conversion/Passes.td@@ -12,7 +12,6 @@ def ConvertTritonToTritonGPU: Pass<"convert-triton-to-tritongpu", "mlir::ModuleO let dependentDialects = ["mlir::arith::ArithmeticDialect", "mlir::math::MathDialect",- "mlir::StandardOpsDialect", // TODO: Does this pass depend on SCF? "mlir::scf::SCFDialect", "mlir::triton::TritonDialect",@@ -41,8 +40,7 @@ def ConvertTritonGPUToLLVM : Pass<"convert-triton-gpu-to-llvm", "mlir::ModuleOp" "mlir::tensor::TensorDialect", "mlir::triton::TritonDialect", "mlir::triton::gpu::TritonGPUDialect",- "mlir::NVVM::NVVMDialect",- "mlir::StandardOpsDialect"];+ "mlir::NVVM::NVVMDialect"]; let options = [ Option<"computeCapability", "compute-capability",diff --git a/include/triton/Dialect/Triton/IR/Dialect.h b/include/triton/Dialect/Triton/IR/Dialect.hindex e8012a51df..15869e262e 100644--- a/include/triton/Dialect/Triton/IR/Dialect.h+++ b/include/triton/Dialect/Triton/IR/Dialect.h@@ -1,14 +1,15 @@ #ifndef TRITON_DIALECT_TRITON_IR_DIALECT_H_ #define TRITON_DIALECT_TRITON_IR_DIALECT_H_ +#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"+#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"+#include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Math/IR/Math.h"-#include "mlir/Dialect/SCF/SCF.h"-#include "mlir/Dialect/StandardOps/IR/Ops.h"+#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/Dialect.h" #include "mlir/Interfaces/ControlFlowInterfaces.h"- #include "triton/Dialect/Triton/IR/Dialect.h.inc" #include "triton/Dialect/Triton/IR/OpsEnums.h.inc" #include "triton/Dialect/Triton/IR/Traits.h"diff --git a/include/triton/Dialect/Triton/IR/TritonDialect.td b/include/triton/Dialect/Triton/IR/TritonDialect.tdindex 07b069e14f..d98ce73884 100644--- a/include/triton/Dialect/Triton/IR/TritonDialect.td+++ b/include/triton/Dialect/Triton/IR/TritonDialect.td@@ -25,12 +25,9 @@ def Triton_Dialect : Dialect { let dependentDialects = [ "arith::ArithmeticDialect", "math::MathDialect",- "StandardOpsDialect", "scf::SCFDialect",-- // Since LLVM 15- // "cf::ControlFlowDialect",- // "func::FuncDialect"+ "cf::ControlFlowDialect",+ "func::FuncDialect" ]; let extraClassDeclaration = [{@@ -38,6 +35,7 @@ def Triton_Dialect : Dialect { }]; let hasConstantMaterializer = 1;+ let useDefaultTypePrinterParser = 1; } include "triton/Dialect/Triton/IR/TritonTypes.td"diff --git a/include/triton/Dialect/Triton/IR/TritonOps.td b/include/triton/Dialect/Triton/IR/TritonOps.tdindex 779e0b648c..0a69211179 100644--- a/include/triton/Dialect/Triton/IR/TritonOps.td+++ b/include/triton/Dialect/Triton/IR/TritonOps.td@@ -141,11 +141,7 @@ def TT_LoadOp : TT_Op<"load", "triton::EvictionPolicy":$evict, "bool":$isVolatile)>, ]; - // let assemblyFormat = "operands attr-dict `:` type($result)";- let parser = [{ return mlir::triton::parseLoadOp(parser, result); }];-- let printer = [{ return mlir::triton::printLoadOp(p, *this); }];-+ let hasCustomAssemblyFormat = 1; let hasCanonicalizer = 1; } @@ -170,11 +166,7 @@ def TT_StoreOp : TT_Op<"store", "triton::EvictionPolicy":$evict)>, ]; - // let assemblyFormat = "operands attr-dict `:` type($value)";- let parser = [{ return mlir::triton::parseStoreOp(parser, result); }];-- let printer = [{ return mlir::triton::printStoreOp(p, *this); }];-+ let hasCustomAssemblyFormat = 1; let hasCanonicalizer = 1; } diff --git a/include/triton/Dialect/Triton/IR/TritonTypes.td b/include/triton/Dialect/Triton/IR/TritonTypes.tdindex 66d2a7b9a9..2fe2fd077d 100644--- a/include/triton/Dialect/Triton/IR/TritonTypes.td+++ b/include/triton/Dialect/Triton/IR/TritonTypes.td@@ -1,6 +1,7 @@ #ifndef TRITON_TYPES #define TRITON_TYPES +include "mlir/IR/AttrTypeBase.td" include "triton/Dialect/Triton/IR/TritonDialect.td" //@@ -58,6 +59,7 @@ def TT_Ptr : TritonTypeDef<"Pointer", "ptr"> { }]> ]; + let hasCustomAssemblyFormat = 1; let skipDefaultBuilders = 1; } def TT_PtrTensor : TensorOf<[TT_Ptr]>;diff --git a/include/triton/Dialect/Triton/Transforms/Passes.td b/include/triton/Dialect/Triton/Transforms/Passes.tdindex 8f77aed774..a25cdc5680 100644--- a/include/triton/Dialect/Triton/Transforms/Passes.td+++ b/include/triton/Dialect/Triton/Transforms/Passes.td@@ -16,8 +16,7 @@ def TritonCombineOps : Pass</*cli-arg*/"triton-combine", /*Op*/"mlir::ModuleOp"> let constructor = "mlir::triton::createCombineOpsPass()"; - let dependentDialects = ["mlir::arith::ArithmeticDialect",- /*SelectOp*/"mlir::StandardOpsDialect"];+ let dependentDialects = ["mlir::arith::ArithmeticDialect"]; } #endifdiff --git a/include/triton/Dialect/TritonGPU/IR/Dialect.h b/include/triton/Dialect/TritonGPU/IR/Dialect.hindex b4c8daec7b..dfc5f53ab1 100644--- a/include/triton/Dialect/TritonGPU/IR/Dialect.h+++ b/include/triton/Dialect/TritonGPU/IR/Dialect.h@@ -1,19 +1,17 @@ #ifndef TRITON_DIALECT_TRITONGPU_IR_DIALECT_H_ #define TRITON_DIALECT_TRITONGPU_IR_DIALECT_H_ -#include "mlir/Dialect/GPU/GPUDialect.h"+#include "mlir/Dialect/GPU/IR/GPUDialect.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/Dialect.h" // TritonGPU depends on Triton #include "triton/Dialect/Triton/IR/Dialect.h"- #include "triton/Dialect/TritonGPU/IR/Dialect.h.inc" #include "triton/Dialect/TritonGPU/IR/Traits.h" #define GET_ATTRDEF_CLASSES-#include "triton/Dialect/Triton/IR/AttrInterfaces.h.inc" #include "triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.h.inc" #define GET_OP_CLASSESdiff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td b/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.tdindex 0242c3cc17..af2aeb03a8 100644--- a/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td+++ b/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td@@ -1,6 +1,7 @@ #ifndef TRITONGPU_ATTRDEFS #define TRITONGPU_ATTRDEFS +include "mlir/IR/AttrTypeBase.td" include "triton/Dialect/TritonGPU/IR/TritonGPUDialect.td" include "triton/Dialect/Triton/IR/TritonInterfaces.td" @@ -136,6 +137,7 @@ A_{3, 2} A_{3, 3} A_{3, 0} A_{3, 1} ... [phase 1] / ]; let extraClassDeclaration = extraBaseClassDeclaration;+ let hasCustomAssemblyFormat = 1; } //===----------------------------------------------------------------------===//@@ -273,6 +275,7 @@ for // ArrayRefParameter<"unsigned">:$sizePerCTA ); + let hasCustomAssemblyFormat = 1; } //===----------------------------------------------------------------------===//@@ -422,6 +425,7 @@ For example, the matrix L corresponding to blockTileSize=[32,16] is: static constexpr int numBitsToHoldMmaV1ID{5}; }]; + let hasCustomAssemblyFormat = 1; } def SliceEncodingAttr : DistributedEncoding<"SliceEncoding"> {@@ -456,6 +460,8 @@ def SliceEncodingAttr : DistributedEncoding<"SliceEncoding"> { template<class T> SmallVector<T> paddedShape(ArrayRef<T> shape) const; }];++ let hasCustomAssemblyFormat = 1; } def DotOperandEncodingAttr : DistributedEncoding<"DotOperandEncoding"> {@@ -492,6 +498,7 @@ section 9.7.13.4.1 for more details. ]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = extraBaseClassDeclaration; } diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.td b/include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.tdindex 87ec1d36c6..6489a721b4 100644--- a/include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.td+++ b/include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.td@@ -30,7 +30,7 @@ def TritonGPU_Dialect : Dialect { } }]; -+ let useDefaultAttributePrinterParser = 1; } #endifdiff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td b/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.tdindex 510f8d0183..7aba11dc75 100644--- a/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td+++ b/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td@@ -59,7 +59,7 @@ def TTG_AsyncCommitGroupOp : TTG_Op<"async_commit_group"> { // This is needed because these ops don't // handle encodings // e.g., https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td#L111-def TTG_CmpIOp : TTG_Op<"cmpi", [NoSideEffect, Elementwise, +def TTG_CmpIOp : TTG_Op<"cmpi", [NoSideEffect, Elementwise, SameOperandsAndResultShape, SameOperandsAndResultEncoding]> { let summary = "integer comparison operation";@@ -73,7 +73,7 @@ def TTG_CmpIOp : TTG_Op<"cmpi", [NoSideEffect, Elementwise, let results = (outs TT_BoolLike:$result); } -def TTG_CmpFOp : TTG_Op<"cmpf", [NoSideEffect, Elementwise, +def TTG_CmpFOp : TTG_Op<"cmpf", [NoSideEffect, Elementwise, SameOperandsAndResultShape, SameOperandsAndResultEncoding]> { let summary = "floating-point comparison operation";@@ -88,8 +88,8 @@ def TTG_CmpFOp : TTG_Op<"cmpf", [NoSideEffect, Elementwise, } // TODO: migrate to arith::SelectOp on LLVM16-def TTG_SelectOp : TTG_Op<"select", [NoSideEffect, Elementwise, - SameOperandsAndResultShape, +def TTG_SelectOp : TTG_Op<"select", [NoSideEffect, Elementwise,+ SameOperandsAndResultShape, SameOperandsAndResultEncoding]> { let summary = "select operation"; @@ -188,10 +188,7 @@ def TTG_InsertSliceAsyncOp : TTG_Op<"insert_slice_async", } }]; - // The custom parser could be replaced with oilist in LLVM-16- let parser = [{ return parseInsertSliceAsyncOp(parser, result); }];-- let printer = [{ return printInsertSliceAsyncOp(p, *this); }];+ let hasCustomAssemblyFormat = 1; } def TTG_AllocTensorOp : TTG_Op<"alloc_tensor", [MemoryEffects<[MemAlloc]>, // Allocate shared memorydiff --git a/lib/Analysis/Alias.cpp b/lib/Analysis/Alias.cppindex a39e4de9aa..208fdd4afc 100644--- a/lib/Analysis/Alias.cpp+++ b/lib/Analysis/Alias.cpp@@ -18,8 +18,9 @@ AliasInfo AliasInfo::join(const AliasInfo &lhs, const AliasInfo &rhs) { return ret; } -ChangeResult SharedMemoryAliasAnalysis::visitOperation(- Operation *op, ArrayRef<LatticeElement<AliasInfo> *> operands) {+void SharedMemoryAliasAnalysis::visitOperation(+ Operation *op, ArrayRef<const dataflow::Lattice<AliasInfo> *> operands,+ ArrayRef<dataflow::Lattice<AliasInfo> *> results) { AliasInfo aliasInfo; bool pessimistic = true; if (maybeSharedAllocationOp(op)) {@@ -44,14 +45,11 @@ ChangeResult SharedMemoryAliasAnalysis::visitOperation( } if (pessimistic) {- return markAllPessimisticFixpoint(op->getResults());+ return markAllPessimisticFixpoint(results); } // Join all lattice elements- ChangeResult result = ChangeResult::NoChange;- for (Value value : op->getResults()) {- result |= getLatticeElement(value).join(aliasInfo);- }- return result;+ for (auto *result : results)+ propagateIfChanged(result, result->join(aliasInfo)); } AliasResult SharedMemoryAliasAnalysis::alias(Value lhs, Value rhs) {diff --git a/lib/Analysis/Allocation.cpp b/lib/Analysis/Allocation.cppindex 712c08c475..b4de8dcd9d 100644--- a/lib/Analysis/Allocation.cpp+++ b/lib/Analysis/Allocation.cpp@@ -1,4 +1,5 @@ #include "triton/Analysis/Allocation.h"+#include "mlir/Analysis/DataFlowFramework.h" #include "mlir/Analysis/Liveness.h" #include "mlir/Analysis/SliceAnalysis.h" #include "mlir/Dialect/Tensor/IR/Tensor.h"@@ -33,10 +34,8 @@ constexpr int kPtrBitWidth = 64; static std::pair<SmallVector<unsigned>, SmallVector<unsigned>> getCvtOrder(const Attribute &srcLayout, const Attribute &dstLayout) {- auto srcBlockedLayout = srcLayout.dyn_cast<BlockedEncodingAttr>(); auto srcMmaLayout = srcLayout.dyn_cast<MmaEncodingAttr>(); auto srcDotLayout = srcLayout.dyn_cast<DotOperandEncodingAttr>();- auto dstBlockedLayout = dstLayout.dyn_cast<BlockedEncodingAttr>(); auto dstMmaLayout = dstLayout.dyn_cast<MmaEncodingAttr>(); auto dstDotLayout = dstLayout.dyn_cast<DotOperandEncodingAttr>(); assert(!(srcMmaLayout && dstMmaLayout) &&@@ -224,14 +223,12 @@ class AllocationAnalysis { } void getValueAlias(Value value, SharedMemoryAliasAnalysis &analysis) {- LatticeElement<AliasInfo> *latticeElement =- analysis.lookupLatticeElement(value);- if (latticeElement) {- auto &info = latticeElement->getValue();- if (!info.getAllocs().empty()) {- for (auto alloc : info.getAllocs()) {- allocation->addAlias(value, alloc);- }+ dataflow::Lattice<AliasInfo> *latticeElement =+ analysis.getLatticeElement(value);+ if (latticeElement && !latticeElement->isUninitialized()) {+ AliasInfo &info = latticeElement->getValue();+ for (auto alloc : info.getAllocs()) {+ allocation->addAlias(value, alloc); } } }@@ -244,14 +241,19 @@ class AllocationAnalysis { getScratchValueSize(op); }); // Get the alias values- SharedMemoryAliasAnalysis aliasAnalysis(operation->getContext());- aliasAnalysis.run(operation);+ std::unique_ptr<DataFlowSolver> solver = createDataFlowSolver();+ SharedMemoryAliasAnalysis *aliasAnalysis =+ solver->load<SharedMemoryAliasAnalysis>();+ if (failed(solver->initializeAndRun(operation))) {+ // TODO: return error instead of bailing out..+ llvm_unreachable("failed to run SharedMemoryAliasAnalysis");+ } operation->walk<WalkOrder::PreOrder>([&](Operation *op) { for (auto operand : op->getOperands()) {- getValueAlias(operand, aliasAnalysis);+ getValueAlias(operand, *aliasAnalysis); } for (auto value : op->getResults()) {- getValueAlias(value, aliasAnalysis);+ getValueAlias(value, *aliasAnalysis); } }); }diff --git a/lib/Analysis/AxisInfo.cpp b/lib/Analysis/AxisInfo.cppindex 0b7142b04d..4af46c3fbb 100644--- a/lib/Analysis/AxisInfo.cpp+++ b/lib/Analysis/AxisInfo.cpp@@ -1,4 +1,4 @@-#include "mlir/Analysis/DataFlowAnalysis.h"+#include "mlir/Analysis/DataFlowFramework.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "llvm/Support/raw_ostream.h" @@ -52,7 +52,7 @@ AxisInfo AxisInfo::getPessimisticValueState(Value value) { BlockArgument blockArg = value.dyn_cast<BlockArgument>(); if (blockArg && blockArg.getOwner()->isEntryBlock()) { Operation *op = blockArg.getOwner()->getParentOp();- if (FuncOp fun = dyn_cast<FuncOp>(op)) {+ if (func::FuncOp fun = dyn_cast<func::FuncOp>(op)) { Attribute attr = fun.getArgAttr(blockArg.getArgNumber(), "tt.divisibility"); if (attr)@@ -136,8 +136,9 @@ class CastOpAxisInfoVisitor final : public AxisInfoVisitorImpl<OpTy> { public: using AxisInfoVisitorImpl<OpTy>::AxisInfoVisitorImpl; - AxisInfo getAxisInfo(OpTy op,- ArrayRef<LatticeElement<AxisInfo> *> operands) override {+ AxisInfo+ getAxisInfo(OpTy op,+ ArrayRef<const dataflow::Lattice<AxisInfo> *> operands) override { return operands[0]->getValue(); } };@@ -147,8 +148,9 @@ class MakeRangeOpAxisInfoVisitor final public: using AxisInfoVisitorImpl<triton::MakeRangeOp>::AxisInfoVisitorImpl; - AxisInfo getAxisInfo(triton::MakeRangeOp op,- ArrayRef<LatticeElement<AxisInfo> *> operands) override {+ AxisInfo+ getAxisInfo(triton::MakeRangeOp op,+ ArrayRef<const dataflow::Lattice<AxisInfo> *> operands) override { auto start = op.start(); auto end = op.end(); return AxisInfo(/*contiguity=*/{end - start},@@ -162,8 +164,9 @@ class ConstantOpAxisInfoVisitor final public: using AxisInfoVisitorImpl<arith::ConstantOp>::AxisInfoVisitorImpl; - AxisInfo getAxisInfo(arith::ConstantOp op,- ArrayRef<LatticeElement<AxisInfo> *> operands) override {+ AxisInfo+ getAxisInfo(arith::ConstantOp op,+ ArrayRef<const dataflow::Lattice<AxisInfo> *> operands) override { auto intAttr = op.getValue().dyn_cast<IntegerAttr>(); auto boolAttr = op.getValue().dyn_cast<BoolAttr>(); if (intAttr || boolAttr) {@@ -416,8 +419,9 @@ class SplatOpAxisInfoVisitor final public: using AxisInfoVisitorImpl<triton::SplatOp>::AxisInfoVisitorImpl; - AxisInfo getAxisInfo(triton::SplatOp op,- ArrayRef<LatticeElement<AxisInfo> *> operands) override {+ AxisInfo+ getAxisInfo(triton::SplatOp op,+ ArrayRef<const dataflow::Lattice<AxisInfo> *> operands) override { Type _retTy = *op->result_type_begin(); TensorType retTy = _retTy.cast<TensorType>(); AxisInfo opInfo = operands[0]->getValue();@@ -439,8 +443,9 @@ class ExpandDimsOpAxisInfoVisitor final public: using AxisInfoVisitorImpl<triton::ExpandDimsOp>::AxisInfoVisitorImpl; - AxisInfo getAxisInfo(triton::ExpandDimsOp op,- ArrayRef<LatticeElement<AxisInfo> *> operands) override {+ AxisInfo+ getAxisInfo(triton::ExpandDimsOp op,+ ArrayRef<const dataflow::Lattice<AxisInfo> *> operands) override { AxisInfo opInfo = operands[0]->getValue(); AxisInfo::DimVectorT contiguity = opInfo.getContiguity(); AxisInfo::DimVectorT divisibility = opInfo.getDivisibility();@@ -458,8 +463,9 @@ class BroadcastOpAxisInfoVisitor final public: using AxisInfoVisitorImpl<triton::BroadcastOp>::AxisInfoVisitorImpl; - AxisInfo getAxisInfo(triton::BroadcastOp op,- ArrayRef<LatticeElement<AxisInfo> *> operands) override {+ AxisInfo+ getAxisInfo(triton::BroadcastOp op,+ ArrayRef<const dataflow::Lattice<AxisInfo> *> operands) override { Type _retTy = *op->result_type_begin(); Type _opTy = *op->operand_type_begin(); TensorType retTy = _retTy.cast<TensorType>();@@ -486,8 +492,9 @@ class CmpOpAxisInfoVisitor final : public AxisInfoVisitorImpl<OpTy> { public: using AxisInfoVisitorImpl<OpTy>::AxisInfoVisitorImpl; - AxisInfo getAxisInfo(OpTy op,- ArrayRef<LatticeElement<AxisInfo> *> operands) override {+ AxisInfo+ getAxisInfo(OpTy op,+ ArrayRef<const dataflow::Lattice<AxisInfo> *> operands) override { auto resTy = op.getResult().getType().template dyn_cast<RankedTensorType>(); if (!resTy) return AxisInfo();@@ -596,8 +603,9 @@ class SelectOpAxisInfoVisitor final : public AxisInfoVisitorImpl<OpTy> { public: using AxisInfoVisitorImpl<OpTy>::AxisInfoVisitorImpl; - AxisInfo getAxisInfo(OpTy op,- ArrayRef<LatticeElement<AxisInfo> *> operands) override {+ AxisInfo+ getAxisInfo(OpTy op,+ ArrayRef<const dataflow::Lattice<AxisInfo> *> operands) override { auto resTy = op.getResult().getType().template dyn_cast<RankedTensorType>(); if (!resTy) return AxisInfo();@@ -757,8 +765,9 @@ class MaxMinOpAxisInfoVisitor final : public AxisInfoVisitorImpl<OpTy> { public: using AxisInfoVisitorImpl<OpTy>::AxisInfoVisitorImpl; - AxisInfo getAxisInfo(OpTy op,- ArrayRef<LatticeElement<AxisInfo> *> operands) override {+ AxisInfo+ getAxisInfo(OpTy op,+ ArrayRef<const dataflow::Lattice<AxisInfo> *> operands) override { auto lhsInfo = operands[0]->getValue(); auto rhsInfo = operands[1]->getValue(); std::optional<int64_t> constantValue;@@ -786,8 +795,8 @@ class MaxMinOpAxisInfoVisitor final : public AxisInfoVisitorImpl<OpTy> { // AxisInfoAnalysis //===----------------------------------------------------------------------===// -AxisInfoAnalysis::AxisInfoAnalysis(MLIRContext *context)- : ForwardDataFlowAnalysis<AxisInfo>(context) {+AxisInfoAnalysis::AxisInfoAnalysis(DataFlowSolver &solver)+ : dataflow::SparseDataFlowAnalysis<dataflow::Lattice<AxisInfo>>(solver) { // UnrealizedConversionCast: // This is needed by TritonGPUToLLVM, to get AxisInfo when the graph is // in the process of a PartialConversion, where UnrealizedConversionCast@@ -819,7 +828,7 @@ AxisInfoAnalysis::AxisInfoAnalysis(MLIRContext *context) visitors.append<LogicalOpAxisInfoVisitor<arith::AndIOp>, LogicalOpAxisInfoVisitor<arith::OrIOp>, LogicalOpAxisInfoVisitor<arith::XOrIOp>>();- visitors.append<SelectOpAxisInfoVisitor<mlir::SelectOp>,+ visitors.append<SelectOpAxisInfoVisitor<mlir::arith::SelectOp>, SelectOpAxisInfoVisitor<triton::gpu::SelectOp>>(); visitors.append<ShLIOpAxisInfoVisitor, ShROpAxisInfoVisitor<arith::ShRUIOp>, ShROpAxisInfoVisitor<arith::ShRSIOp>>();@@ -829,11 +838,12 @@ AxisInfoAnalysis::AxisInfoAnalysis(MLIRContext *context) MaxMinOpAxisInfoVisitor<arith::MinUIOp>>(); } -ChangeResult AxisInfoAnalysis::visitOperation(- Operation *op, ArrayRef<LatticeElement<AxisInfo> *> operands) {+void AxisInfoAnalysis::visitOperation(+ Operation *op, ArrayRef<const dataflow::Lattice<AxisInfo> *> operands,+ ArrayRef<dataflow::Lattice<AxisInfo> *> results) { AxisInfo curr = visitors.apply(op, operands); if (curr.getRank() == 0) {- return markAllPessimisticFixpoint(op->getResults());+ return markAllPessimisticFixpoint(results); } // override with hint auto newContiguity = curr.getContiguity();@@ -854,11 +864,8 @@ ChangeResult AxisInfoAnalysis::visitOperation( curr = mlir::AxisInfo(newContiguity, newDivisibility, newConstancy, curr.getConstantValue()); // join all lattice elements- ChangeResult result = ChangeResult::NoChange;- for (Value value : op->getResults()) {- result |= getLatticeElement(value).join(curr);- }- return result;+ for (auto *result : results)+ propagateIfChanged(result, result->join(curr)); } unsigned AxisInfoAnalysis::getPtrContiguity(Value ptr) {@@ -884,7 +891,10 @@ unsigned AxisInfoAnalysis::getPtrAlignment(Value ptr) { auto tensorTy = ptr.getType().dyn_cast<RankedTensorType>(); if (!tensorTy) return 1;- auto axisInfo = lookupLatticeElement(ptr)->getValue();+ dataflow::Lattice<AxisInfo> *latticeElement = getLatticeElement(ptr);+ if (!latticeElement || latticeElement->isUninitialized())+ return 1;+ auto axisInfo = latticeElement->getValue(); auto layout = tensorTy.getEncoding(); auto order = triton::gpu::getOrder(layout); auto maxMultipleBytes = axisInfo.getDivisibility(order[0]);@@ -900,8 +910,11 @@ unsigned AxisInfoAnalysis::getMaskAlignment(Value mask) { auto tensorTy = mask.getType().dyn_cast<RankedTensorType>(); if (!tensorTy) return 1;+ dataflow::Lattice<AxisInfo> *latticeElement = getLatticeElement(mask);+ if (!latticeElement || latticeElement->isUninitialized())+ return 1;+ auto maskAxis = latticeElement->getValue(); auto maskOrder = triton::gpu::getOrder(tensorTy.getEncoding());- auto maskAxis = lookupLatticeElement(mask)->getValue(); auto alignment = std::max<unsigned>(maskAxis.getConstancy(maskOrder[0]), 1); return alignment; }diff --git a/lib/Analysis/CMakeLists.txt b/lib/Analysis/CMakeLists.txtindex afbc692510..1f761f845c 100644--- a/lib/Analysis/CMakeLists.txt+++ b/lib/Analysis/CMakeLists.txt@@ -8,7 +8,7 @@ add_mlir_library(TritonAnalysis DEPENDS TritonTableGen TritonGPUAttrDefsIncGen- + LINK_LIBS PUBLIC MLIRAnalysis )diff --git a/lib/Analysis/Membar.cpp b/lib/Analysis/Membar.cppindex acc885e827..910274b2ac 100644--- a/lib/Analysis/Membar.cpp+++ b/lib/Analysis/Membar.cpp@@ -2,7 +2,7 @@ #include "triton/Analysis/Alias.h" #include "triton/Dialect/TritonGPU/IR/Dialect.h" -#include "mlir/Dialect/GPU/GPUDialect.h"+#include "mlir/Dialect/GPU/IR/GPUDialect.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" namespace mlir {diff --git a/lib/Analysis/Utility.cpp b/lib/Analysis/Utility.cppindex d9e917e731..6ea52df272 100644--- a/lib/Analysis/Utility.cpp+++ b/lib/Analysis/Utility.cpp@@ -1,5 +1,8 @@ #include "triton/Analysis/Utility.h"+#include "mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h"+#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h" #include "mlir/IR/Dialect.h"+#include "mlir/IR/Matchers.h" #include "triton/Dialect/Triton/IR/Dialect.h" #include "triton/Dialect/TritonGPU/IR/Dialect.h" #include <deque>@@ -325,4 +328,55 @@ SetVector<Operation *> multiRootGetSlice(Operation *op, return multiRootTopologicalSort(slice); } +namespace {+// Copied from TestDeadCodeAnalysis.cpp, because some dead code analysis+// interacts with constant propagation, but SparseConstantPropagation+// doesn't seem to be sufficient.+struct ConstantAnalysis : public DataFlowAnalysis {+ using DataFlowAnalysis::DataFlowAnalysis;++ LogicalResult initialize(Operation *top) override {+ WalkResult result = top->walk([&](Operation *op) {+ if (failed(visit(op)))+ return WalkResult::interrupt();+ return WalkResult::advance();+ });+ return success(!result.wasInterrupted());+ }++ LogicalResult visit(ProgramPoint point) override {+ Operation *op = point.get<Operation *>();+ Attribute value;+ if (matchPattern(op, m_Constant(&value))) {+ auto *constant = getOrCreate<dataflow::Lattice<dataflow::ConstantValue>>(+ op->getResult(0));+ propagateIfChanged(constant, constant->join(dataflow::ConstantValue(+ value, op->getDialect())));+ return success();+ }+ setAllToUnknownConstants(op->getResults());+ for (Region ®ion : op->getRegions())+ setAllToUnknownConstants(region.getArguments());+ return success();+ }++ /// Set all given values as not constants.+ void setAllToUnknownConstants(ValueRange values) {+ dataflow::ConstantValue unknownConstant(nullptr, nullptr);+ for (Value value : values) {+ auto *constant =+ getOrCreate<dataflow::Lattice<dataflow::ConstantValue>>(value);+ propagateIfChanged(constant, constant->join(unknownConstant));+ }+ }+};+} // namespace++std::unique_ptr<DataFlowSolver> createDataFlowSolver() {+ auto solver = std::make_unique<DataFlowSolver>();+ solver->load<dataflow::DeadCodeAnalysis>();+ solver->load<ConstantAnalysis>();+ return solver;+}+ } // namespace mlirdiff --git a/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cppindex 6a46265bd7..e352eb3698 100644--- a/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp+++ b/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp@@ -159,9 +159,6 @@ struct ConvertLayoutOpConversion Value smemBase) const { auto accumNumCTAsEachRep = product<unsigned>(numCTAsEachRep); auto layout = type.getEncoding();- auto blockedLayout = layout.dyn_cast<BlockedEncodingAttr>();- auto sliceLayout = layout.dyn_cast<SliceEncodingAttr>();- auto mmaLayout = layout.dyn_cast<MmaEncodingAttr>(); auto rank = type.getRank(); auto sizePerThread = getSizePerThread(layout); auto accumSizePerThread = product<unsigned>(sizePerThread);diff --git a/lib/Conversion/TritonGPUToLLVM/DotOpHelpers.h b/lib/Conversion/TritonGPUToLLVM/DotOpHelpers.hindex 4b89965aa9..1d9e00519b 100644--- a/lib/Conversion/TritonGPUToLLVM/DotOpHelpers.h+++ b/lib/Conversion/TritonGPUToLLVM/DotOpHelpers.h@@ -7,10 +7,8 @@ #include "mlir/Conversion/LLVMCommon/LoweringOptions.h" #include "mlir/Conversion/LLVMCommon/Pattern.h" #include "mlir/Conversion/MathToLLVM/MathToLLVM.h"-#include "mlir/Conversion/SCFToStandard/SCFToStandard.h"-#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h" #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"-#include "mlir/Dialect/GPU/GPUDialect.h"+#include "mlir/Dialect/GPU/IR/GPUDialect.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/Matchers.h"@@ -422,9 +420,9 @@ struct MMA16816ConversionHelper { MMA16816ConversionHelper(Type dotOperand, MmaEncodingAttr mmaLayout, Value thread, ConversionPatternRewriter &rewriter, TypeConverter *typeConverter, Location loc)- : mmaLayout(mmaLayout), thread(thread), helper(mmaLayout),- rewriter(rewriter), typeConverter(typeConverter), loc(loc),- ctx(mmaLayout.getContext()), wpt(mmaLayout.getWarpsPerCTA()) {+ : mmaLayout(mmaLayout), wpt(mmaLayout.getWarpsPerCTA()), thread(thread),+ helper(mmaLayout), rewriter(rewriter), typeConverter(typeConverter),+ loc(loc), ctx(mmaLayout.getContext()) { helper.deduceMmaType(dotOperand); Value _32 = i32_val(32);diff --git a/lib/Conversion/TritonGPUToLLVM/DotOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/DotOpToLLVM.cppindex 0f8070ca9f..e4bd47c411 100644--- a/lib/Conversion/TritonGPUToLLVM/DotOpToLLVM.cpp+++ b/lib/Conversion/TritonGPUToLLVM/DotOpToLLVM.cpp@@ -115,8 +115,6 @@ struct DotOpConversion : public ConvertTritonGPUOpToLLVMPattern<triton::DotOp> { auto DTensorTy = D.getType().cast<RankedTensorType>(); auto AShape = ATensorTy.getShape(); auto BShape = BTensorTy.getShape();- auto DShape = DTensorTy.getShape();- auto wpt = mmaLayout.getWarpsPerCTA(); bool isARow = ALayout.getIsMMAv1Row().cast<BoolAttr>().getValue(); bool isBRow = BLayout.getIsMMAv1Row().cast<BoolAttr>().getValue();@@ -221,7 +219,6 @@ struct DotOpConversion : public ConvertTritonGPUOpToLLVMPattern<triton::DotOp> { ConversionPatternRewriter &rewriter) const { auto *ctx = rewriter.getContext(); auto loc = op.getLoc();- auto threadId = getThreadId(rewriter, loc); auto A = op.a(); auto B = op.b();@@ -230,12 +227,10 @@ struct DotOpConversion : public ConvertTritonGPUOpToLLVMPattern<triton::DotOp> { auto aTensorTy = A.getType().cast<RankedTensorType>(); auto bTensorTy = B.getType().cast<RankedTensorType>();- auto cTensorTy = C.getType().cast<RankedTensorType>(); auto dTensorTy = D.getType().cast<RankedTensorType>(); auto aShape = aTensorTy.getShape(); auto bShape = bTensorTy.getShape();- auto cShape = cTensorTy.getShape(); BlockedEncodingAttr dLayout = dTensorTy.getEncoding().cast<BlockedEncodingAttr>();diff --git a/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cppindex deb71b9597..0b9e67674b 100644--- a/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp+++ b/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp@@ -61,7 +61,6 @@ struct FpToFpOpConversion convertFp16x4ToFp8x4(Location loc, ConversionPatternRewriter &rewriter, const Value &v0, const Value &v1, const Value &v2, const Value &v3) {- auto ctx = rewriter.getContext(); auto fp16x2VecTy = vec_ty(f16_ty, 2); Value fp16x2Vec0 = undef(fp16x2VecTy); Value fp16x2Vec1 = undef(fp16x2VecTy);@@ -153,7 +152,6 @@ struct FpToFpOpConversion convertBf16x4ToFp8x4(Location loc, ConversionPatternRewriter &rewriter, const Value &v0, const Value &v1, const Value &v2, const Value &v3) {- auto ctx = rewriter.getContext(); auto bf16x2VecTy = vec_ty(i16_ty, 2); Value bf16x2Vec0 = undef(bf16x2VecTy); Value bf16x2Vec1 = undef(bf16x2VecTy);diff --git a/lib/Conversion/TritonGPUToLLVM/LoadStoreOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/LoadStoreOpToLLVM.cppindex 9a8b4702bc..bae675f0cb 100644--- a/lib/Conversion/TritonGPUToLLVM/LoadStoreOpToLLVM.cpp+++ b/lib/Conversion/TritonGPUToLLVM/LoadStoreOpToLLVM.cpp@@ -109,7 +109,8 @@ struct LoadOpConversion DenseElementsAttr constAttr; int64_t splatVal = 0; if (other && valueElemTy.isa<IntegerType>() &&- matchPattern(other, m_Constant(&constAttr)) && constAttr.isSplat()) {+ matchPattern(other, m_Constant(&constAttr)) && constAttr.isSplat() &&+ constAttr.getElementType().isa<IntegerType>()) { otherIsSplatConstInt = true; splatVal = constAttr.getSplatValue<APInt>().getSExtValue(); }@@ -333,7 +334,6 @@ struct StoreOpConversion elem = rewriter.create<LLVM::SExtOp>(loc, type::i8Ty(ctx), elem); elem = bitcast(elem, valueElemTy); - Type u32Ty = typeConverter->convertType(type::u32Ty(ctx)); llWord = insert_element(wordTy, llWord, elem, i32_val(elemIdx)); } llWord = bitcast(llWord, valArgTy);@@ -387,7 +387,6 @@ struct AtomicCASOpConversion ConversionPatternRewriter &rewriter) const override { auto loc = op.getLoc(); MLIRContext *ctx = rewriter.getContext();- Value ptr = op.ptr(); Value llPtr = adaptor.ptr(); Value llCmp = adaptor.cmp();diff --git a/lib/Conversion/TritonGPUToLLVM/ReduceOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/ReduceOpToLLVM.cppindex 69abd889be..1c973dc196 100644--- a/lib/Conversion/TritonGPUToLLVM/ReduceOpToLLVM.cpp+++ b/lib/Conversion/TritonGPUToLLVM/ReduceOpToLLVM.cpp@@ -286,7 +286,6 @@ struct ReduceOpConversion auto srcTy = op.operand().getType().cast<RankedTensorType>(); auto srcLayout = srcTy.getEncoding(); auto srcShape = srcTy.getShape();- auto srcRank = srcTy.getRank(); auto order = getOrder(srcLayout); auto threadsPerWarp = triton::gpu::getThreadsPerWarp(srcLayout);@@ -351,7 +350,6 @@ struct ReduceOpConversion Value zero = i32_val(0); Value laneZero = icmp_eq(laneIdAxis, zero);- Value warpZero = icmp_eq(warpIdAxis, zero); for (auto it : accs) { const SmallVector<unsigned> &key = it.first;diff --git a/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVM.cppindex 5b77150b1a..78cfa076bd 100644--- a/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVM.cpp+++ b/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVM.cpp@@ -11,11 +11,11 @@ using ::mlir::LLVM::getStructFromElements; using ::mlir::triton::gpu::getElemsPerThread; using ::mlir::triton::gpu::SharedEncodingAttr; -struct ReturnOpConversion : public ConvertOpToLLVMPattern<::mlir::ReturnOp> {- using ConvertOpToLLVMPattern<ReturnOp>::ConvertOpToLLVMPattern;+struct ReturnOpConversion : public ConvertOpToLLVMPattern<func::ReturnOp> {+ using ConvertOpToLLVMPattern<func::ReturnOp>::ConvertOpToLLVMPattern; LogicalResult- matchAndRewrite(ReturnOp op, OpAdaptor adaptor,+ matchAndRewrite(func::ReturnOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { unsigned numArguments = op.getNumOperands(); @@ -476,7 +476,6 @@ struct ExtractSliceOpConversion auto llvmElemTy = getTypeConverter()->convertType(srcTy.getElementType()); auto elemPtrTy = ptr_ty(llvmElemTy, 3);- auto resTy = op.getType().dyn_cast<RankedTensorType>(); smemObj = SharedMemoryObject(gep(elemPtrTy, smemObj.base, offset), strideVals, offsetVals); auto retVal = getStructFromSharedMemoryObject(loc, smemObj, rewriter);diff --git a/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVMBase.h b/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVMBase.hindex bb10d5b24a..00e399f848 100644--- a/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVMBase.h+++ b/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVMBase.h@@ -4,6 +4,7 @@ // TODO: refactor so that it doesn't fail if Allocation.h // is included after utility.h (due to conflict in `store` macro // and <atomic>+#include "mlir/Dialect/Func/IR/FuncOps.h" #include "triton/Analysis/Allocation.h" //@@ -39,15 +40,15 @@ void vprintf_array(Value thread, ArrayRef<Value> arr, std::string info, // TODO(Superjomn): remove the code when MLIR v15.0 is included. // All the rights are reserved by the LLVM community. -struct FuncOpConversionBase : public ConvertOpToLLVMPattern<FuncOp> {+struct FuncOpConversionBase : public ConvertOpToLLVMPattern<func::FuncOp> { private: /// Only retain those attributes that are not constructed by /// `LLVMFuncOp::build`. If `filterArgAttrs` is set, also filter out argument /// attributes.- static void filterFuncAttributes(ArrayRef<NamedAttribute> attrs,- bool filterArgAttrs,+ static void filterFuncAttributes(func::FuncOp op, bool filterArgAttrs, SmallVectorImpl<NamedAttribute> &result) {- for (const auto &attr : attrs) {++ for (const auto &attr : op->getAttrs()) { if (attr.getName() == SymbolTable::getSymbolAttrName() || attr.getName() == FunctionOpInterface::getTypeAttrName() || attr.getName() == "std.varargs" ||@@ -65,27 +66,27 @@ struct FuncOpConversionBase : public ConvertOpToLLVMPattern<FuncOp> { } protected:- using ConvertOpToLLVMPattern<FuncOp>::ConvertOpToLLVMPattern;+ using ConvertOpToLLVMPattern<func::FuncOp>::ConvertOpToLLVMPattern; // Convert input FuncOp to LLVMFuncOp by using the LLVMTypeConverter provided // to this legalization pattern. LLVM::LLVMFuncOp- convertFuncOpToLLVMFuncOp(FuncOp funcOp,+ convertFuncOpToLLVMFuncOp(func::FuncOp funcOp, ConversionPatternRewriter &rewriter) const { // Convert the original function arguments. They are converted using the // LLVMTypeConverter provided to this legalization pattern. auto varargsAttr = funcOp->getAttrOfType<BoolAttr>("func.varargs"); TypeConverter::SignatureConversion result(funcOp.getNumArguments()); auto llvmType = getTypeConverter()->convertFunctionSignature(- funcOp.getType(), varargsAttr && varargsAttr.getValue(), result);+ funcOp.getFunctionType(), varargsAttr && varargsAttr.getValue(),+ result); if (!llvmType) return nullptr; // Propagate argument/result attributes to all converted arguments/result // obtained after converting a given original argument/result. SmallVector<NamedAttribute, 4> attributes;- filterFuncAttributes(funcOp->getAttrs(), /*filterArgAttrs=*/true,- attributes);+ filterFuncAttributes(funcOp, /*filterArgAttrs=*/true, attributes); if (ArrayAttr resAttrDicts = funcOp.getAllResultAttrs()) { assert(!resAttrDicts.empty() && "expected array to be non-empty"); auto newResAttrDicts =@@ -131,7 +132,7 @@ struct FuncOpConversionBase : public ConvertOpToLLVMPattern<FuncOp> { } auto newFuncOp = rewriter.create<LLVM::LLVMFuncOp>( funcOp.getLoc(), funcOp.getName(), llvmType, linkage,- /*dsoLocal*/ false, attributes);+ /*dsoLocal*/ false, LLVM::CConv::C, attributes); rewriter.inlineRegionBefore(funcOp.getBody(), newFuncOp.getBody(), newFuncOp.end()); if (failed(rewriter.convertRegionTypes(&newFuncOp.getBody(), *typeConverter,@@ -191,8 +192,8 @@ class ConvertTritonGPUOpToLLVMPatternBase { const Allocation *allocation, Value smem, IndexCacheInfo indexCacheInfo)- : converter(&typeConverter), indexCacheInfo(indexCacheInfo),- allocation(allocation), smem(smem) {}+ : converter(&typeConverter), allocation(allocation), smem(smem),+ indexCacheInfo(indexCacheInfo) {} LLVMTypeConverter *getTypeConverter() const { return converter; } @@ -861,7 +862,6 @@ class ConvertTritonGPUOpToLLVMPatternBase { ArrayRef<int64_t> shape) const { auto parent = sliceLayout.getParent(); unsigned dim = sliceLayout.getDim();- size_t rank = shape.size(); auto parentIndices = emitIndices(loc, rewriter, parent, sliceLayout.paddedShape(shape)); unsigned numIndices = parentIndices.size();diff --git a/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVMPass.cpp b/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVMPass.cppindex ff1af09835..6f66af4e34 100644--- a/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVMPass.cpp+++ b/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVMPass.cpp@@ -1,10 +1,11 @@ #include "triton/Conversion/TritonGPUToLLVM/TritonGPUToLLVMPass.h" +#include "mlir/Analysis/DataFlowFramework.h" #include "mlir/Conversion/ArithmeticToLLVM/ArithmeticToLLVM.h"+#include "mlir/Conversion/ControlFlowToLLVM//ControlFlowToLLVM.h" #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h" #include "mlir/Conversion/MathToLLVM/MathToLLVM.h"-#include "mlir/Conversion/SCFToStandard/SCFToStandard.h"-#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"+#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "mlir/Dialect/LLVMIR/NVVMDialect.h" #include "mlir/Pass/Pass.h"@@ -40,7 +41,6 @@ class TritonLLVMConversionTarget : public ConversionTarget { addIllegalDialect<triton::TritonDialect>(); addIllegalDialect<triton::gpu::TritonGPUDialect>(); addIllegalDialect<mlir::gpu::GPUDialect>();- addIllegalDialect<mlir::StandardOpsDialect>(); addLegalOp<mlir::UnrealizedConversionCastOp>(); } };@@ -51,7 +51,7 @@ class TritonLLVMFunctionConversionTarget : public ConversionTarget { : ConversionTarget(ctx) { addLegalDialect<LLVM::LLVMDialect>(); addLegalDialect<NVVM::NVVMDialect>();- addIllegalOp<mlir::FuncOp>();+ addIllegalOp<mlir::func::FuncOp>(); addLegalOp<mlir::UnrealizedConversionCastOp>(); } };@@ -69,7 +69,7 @@ struct FuncOpConversion : public FuncOpConversionBase { : FuncOpConversionBase(converter, benefit), numWarps(numWarps) {} LogicalResult- matchAndRewrite(FuncOp funcOp, OpAdaptor adaptor,+ matchAndRewrite(func::FuncOp funcOp, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { auto newFuncOp = convertFuncOpToLLVMFuncOp(funcOp, rewriter); if (!newFuncOp)@@ -133,7 +133,8 @@ class ConvertTritonGPUToLLVM decomposeBlockedToDotOperand(mod); // Step 2- decomposeInsertSliceAsyncOp(mod);+ if (failed(decomposeInsertSliceAsyncOp(mod)))+ return signalPassFailure(); // Step 3 Allocation allocation(mod);@@ -142,7 +143,7 @@ class ConvertTritonGPUToLLVM // Step 4 RewritePatternSet scf_patterns(context);- mlir::populateLoopToStdConversionPatterns(scf_patterns);+ mlir::populateSCFToControlFlowConversionPatterns(scf_patterns); mlir::ConversionTarget scf_target(*context); scf_target.addIllegalOp<scf::ForOp, scf::IfOp, scf::ParallelOp, scf::WhileOp, scf::ExecuteRegionOp>();@@ -159,8 +160,10 @@ class ConvertTritonGPUToLLVM return signalPassFailure(); // Step 6 - get axis and shared memory info- AxisInfoAnalysis axisInfoAnalysis(mod.getContext());- axisInfoAnalysis.run(mod);+ std::unique_ptr<DataFlowSolver> solver = createDataFlowSolver();+ AxisInfoAnalysis *axisInfoAnalysis = solver->load<AxisInfoAnalysis>();+ if (failed(solver->initializeAndRun(mod)))+ return signalPassFailure(); initSharedMemory(allocation.getSharedMemorySize(), typeConverter); mod->setAttr("triton_gpu.shared", mlir::IntegerAttr::get(mlir::IntegerType::get(context, 32),@@ -178,38 +181,39 @@ class ConvertTritonGPUToLLVM // Normal conversions populateTritonGPUToLLVMPatterns(typeConverter, patterns, numWarps,- axisInfoAnalysis, &allocation, smem,+ *axisInfoAnalysis, &allocation, smem, indexCacheInfo, /*benefit=*/10); // ConvertLayoutOp populateConvertLayoutOpToLLVMPatterns(typeConverter, patterns, numWarps,- axisInfoAnalysis, &allocation, smem,+ *axisInfoAnalysis, &allocation, smem, indexCacheInfo, /*benefit=*/10); // DotOp populateDotOpToLLVMPatterns(typeConverter, patterns, numWarps,- axisInfoAnalysis, &allocation, smem,+ *axisInfoAnalysis, &allocation, smem, /*benefit=*/10); // ElementwiseOp populateElementwiseOpToLLVMPatterns(typeConverter, patterns, numWarps,- axisInfoAnalysis, &allocation, smem,+ *axisInfoAnalysis, &allocation, smem, /*benefit=*/10); // LoadStoreOp populateLoadStoreOpToLLVMPatterns(typeConverter, patterns, numWarps,- axisInfoAnalysis, &allocation, smem,+ *axisInfoAnalysis, &allocation, smem, indexCacheInfo, /*benefit=*/10); // ReduceOp populateReduceOpToLLVMPatterns(typeConverter, patterns, numWarps,- axisInfoAnalysis, &allocation, smem,+ *axisInfoAnalysis, &allocation, smem, indexCacheInfo, /*benefit=*/10); // ViewOp populateViewOpToLLVMPatterns(typeConverter, patterns, numWarps,- axisInfoAnalysis, &allocation, smem,+ *axisInfoAnalysis, &allocation, smem, /*benefit=*/10); // Add arith/math's patterns to help convert scalar expression to LLVM. mlir::arith::populateArithmeticToLLVMConversionPatterns(typeConverter, patterns); mlir::populateMathToLLVMConversionPatterns(typeConverter, patterns);- mlir::populateStdToLLVMConversionPatterns(typeConverter, patterns);+ mlir::cf::populateControlFlowToLLVMConversionPatterns(typeConverter,+ patterns); mlir::populateGpuToNVVMConversionPatterns(typeConverter, patterns); if (failed(applyPartialConversion(mod, target, std::move(patterns))))@@ -306,9 +310,11 @@ class ConvertTritonGPUToLLVM }); } - void decomposeInsertSliceAsyncOp(ModuleOp mod) const {- AxisInfoAnalysis axisInfoAnalysis(mod.getContext());- axisInfoAnalysis.run(mod);+ LogicalResult decomposeInsertSliceAsyncOp(ModuleOp mod) const {+ std::unique_ptr<DataFlowSolver> solver = createDataFlowSolver();+ AxisInfoAnalysis *axisInfoAnalysis = solver->load<AxisInfoAnalysis>();+ if (failed(solver->initializeAndRun(mod)))+ return failure(); // TODO(Keren): This is a hacky knob that may cause performance regression // when decomposition has been performed. We should remove this knob once we // have thorough analysis on async wait. Currently, we decompose@@ -342,7 +348,7 @@ class ConvertTritonGPUToLLVM auto resSharedLayout = dstTy.getEncoding().dyn_cast<triton::gpu::SharedEncodingAttr>(); auto resElemTy = dstTy.getElementType();- unsigned inVec = axisInfoAnalysis.getPtrContiguity(src);+ unsigned inVec = axisInfoAnalysis->getPtrContiguity(src); unsigned outVec = resSharedLayout.getVec(); unsigned minVec = std::min(outVec, inVec); auto maxBitWidth =@@ -400,11 +406,11 @@ class ConvertTritonGPUToLLVM } else if (decomposed) { // Wait for all previous async ops OpBuilder builder(asyncWaitOp);- auto newAsyncWaitOp =- builder.create<triton::gpu::AsyncWaitOp>(asyncWaitOp.getLoc(), 0);+ builder.create<triton::gpu::AsyncWaitOp>(asyncWaitOp.getLoc(), 0); asyncWaitOp.erase(); } });+ return success(); } }; diff --git a/lib/Conversion/TritonGPUToLLVM/Utility.h b/lib/Conversion/TritonGPUToLLVM/Utility.hindex d35dac28c5..11976908cf 100644--- a/lib/Conversion/TritonGPUToLLVM/Utility.h+++ b/lib/Conversion/TritonGPUToLLVM/Utility.h@@ -220,10 +220,7 @@ struct SharedMemoryObject { ConversionPatternRewriter &rewriter) : base(base) { strides = getStridesFromShapeAndOrder(shape, order, loc, rewriter);-- for (auto idx : order) {- offsets.emplace_back(i32_val(0));- }+ offsets.append(order.size(), i32_val(0)); } SmallVector<Value> getElems() const {diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cppindex fe42202c34..5f230f787f 100644--- a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp+++ b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp@@ -1,10 +1,10 @@ #include "triton/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.h" #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"-#include "mlir/Dialect/GPU/GPUDialect.h"+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"+#include "mlir/Dialect/GPU/IR/GPUDialect.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "mlir/Dialect/LLVMIR/NVVMDialect.h"-#include "mlir/Dialect/StandardOps/IR/Ops.h" #include "mlir/Pass/Pass.h" #include "mlir/Transforms/DialectConversion.h" #include "triton/Dialect/Triton/IR/Dialect.h"@@ -59,10 +59,13 @@ class ArithConstantPattern : public OpConversionPattern<arith::ConstantOp> { Type retType = getTypeConverter()->convertType(op.getType()); auto value = adaptor.getValue().dyn_cast<DenseElementsAttr>(); assert(value);- rewriter.replaceOpWithNewOp<arith::ConstantOp>(- op, retType,- value.reshape(retType) // This is a hack. We just want to add encoding- );+ if (value.getElementType().isInteger(1) && value.isSplat())+ // Workaround until https://reviews.llvm.org/D133743 is included.+ value = DenseElementsAttr::get(retType, value.getSplatValue<bool>());+ else+ // This is a hack. We just want to add encoding+ value = value.reshape(retType);+ rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, retType, value); return success(); } };@@ -127,12 +130,12 @@ void populateArithmeticPatternsAndLegality( } // this shouldn't exist if mlir's SelectOp checked encodings properly-class StdSelectPattern : public OpConversionPattern<SelectOp> {+class StdSelectPattern : public OpConversionPattern<arith::SelectOp> { public:- using OpConversionPattern<SelectOp>::OpConversionPattern;+ using OpConversionPattern<arith::SelectOp>::OpConversionPattern; LogicalResult- matchAndRewrite(SelectOp op, typename SelectOp::Adaptor adaptor,+ matchAndRewrite(arith::SelectOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { Type retType = this->getTypeConverter()->convertType(op.getType()); rewriter.replaceOpWithNewOp<triton::gpu::SelectOp>(@@ -148,8 +151,8 @@ void populateStdPatternsAndLegality(TritonGPUTypeConverter &typeConverter, MLIRContext *context = patterns.getContext(); // Rewrite rule patterns.add<StdSelectPattern>(typeConverter, context);- target.addLegalOp<ReturnOp>(); // this is ok because all functions are inlined- // by the frontend+ target.addLegalOp<func::ReturnOp>(); // this is ok because all functions are+ // inlined by the frontend } void populateMathPatternsAndLegality(TritonGPUTypeConverter &typeConverter,@@ -455,18 +458,19 @@ struct TritonPrintfPattern : public OpConversionPattern<triton::PrintfOp> { void populateTritonPatterns(TritonGPUTypeConverter &typeConverter, RewritePatternSet &patterns) { MLIRContext *context = patterns.getContext();- patterns.add< // TODO: view should have custom pattern that views the layout- TritonGenericPattern<triton::ViewOp>,- TritonGenericPattern<triton::BitcastOp>,- TritonGenericPattern<triton::FpToFpOp>,- TritonGenericPattern<triton::IntToPtrOp>,- TritonGenericPattern<triton::PtrToIntOp>,- TritonGenericPattern<triton::SplatOp>, TritonBroadcastPattern,- TritonGenericPattern<triton::AddPtrOp>, TritonCatPattern,- TritonReducePattern, TritonTransPattern, TritonExpandDimsPattern,- TritonMakeRangePattern, TritonDotPattern, TritonLoadPattern,- TritonStorePattern, TritonExtElemwisePattern, TritonPrintfPattern,- TritonAtomicRMWPattern>(typeConverter, context);+ patterns+ .insert< // TODO: view should have custom pattern that views the layout+ TritonGenericPattern<triton::ViewOp>,+ TritonGenericPattern<triton::BitcastOp>,+ TritonGenericPattern<triton::FpToFpOp>,+ TritonGenericPattern<triton::IntToPtrOp>,+ TritonGenericPattern<triton::PtrToIntOp>,+ TritonGenericPattern<triton::SplatOp>, TritonBroadcastPattern,+ TritonGenericPattern<triton::AddPtrOp>, TritonCatPattern,+ TritonReducePattern, TritonTransPattern, TritonExpandDimsPattern,+ TritonMakeRangePattern, TritonDotPattern, TritonLoadPattern,+ TritonStorePattern, TritonExtElemwisePattern, TritonPrintfPattern,+ TritonAtomicRMWPattern>(typeConverter, context); } //@@ -623,29 +627,28 @@ void populateSCFPatterns(TritonGPUTypeConverter &typeConverter, // CF -class CFBranchPattern : public OpConversionPattern<BranchOp> {+class CFBranchPattern : public OpConversionPattern<cf::BranchOp> { public:- using OpConversionPattern<BranchOp>::OpConversionPattern;+ using OpConversionPattern<cf::BranchOp>::OpConversionPattern; LogicalResult- matchAndRewrite(BranchOp op, BranchOp::Adaptor adaptor,+ matchAndRewrite(cf::BranchOp op, cf::BranchOp::Adaptor adaptor, ConversionPatternRewriter &rewriter) const override {- auto converter = getTypeConverter();- auto newOp = rewriter.replaceOpWithNewOp<BranchOp>(op, op.getSuccessor(),- adaptor.getOperands());+ auto newOp = rewriter.replaceOpWithNewOp<cf::BranchOp>(+ op, op.getSuccessor(), adaptor.getOperands()); return success(); } }; -class CFCondBranchPattern : public OpConversionPattern<CondBranchOp> {+class CFCondBranchPattern : public OpConversionPattern<cf::CondBranchOp> { public:- using OpConversionPattern<CondBranchOp>::OpConversionPattern;+ using OpConversionPattern<cf::CondBranchOp>::OpConversionPattern; LogicalResult- matchAndRewrite(CondBranchOp op, CondBranchOp::Adaptor adaptor,+ matchAndRewrite(cf::CondBranchOp op, cf::CondBranchOp::Adaptor adaptor, ConversionPatternRewriter &rewriter) const override { auto converter = getTypeConverter();- auto newOp = rewriter.replaceOpWithNewOp<CondBranchOp>(+ auto newOp = rewriter.replaceOpWithNewOp<cf::CondBranchOp>( op, adaptor.getCondition(), op.getTrueDest(), adaptor.getTrueDestOperands(), op.getFalseDest(), adaptor.getFalseDestOperands());diff --git a/lib/Dialect/Triton/IR/CMakeLists.txt b/lib/Dialect/Triton/IR/CMakeLists.txtindex 2d679b21fd..705554ba6b 100644--- a/lib/Dialect/Triton/IR/CMakeLists.txt+++ b/lib/Dialect/Triton/IR/CMakeLists.txt@@ -10,11 +10,7 @@ add_mlir_dialect_library(TritonIR LINK_LIBS PUBLIC MLIRIR- MLIRArithmetic- MLIRSCF-- # Since LLVM 15- # MLIRFunc- # else- MLIRStandard+ MLIRArithmeticDialect+ MLIRSCFDialect+ MLIRFuncDialect )diff --git a/lib/Dialect/Triton/IR/Ops.cpp b/lib/Dialect/Triton/IR/Ops.cppindex 3aadbfa0c0..86570359c5 100644--- a/lib/Dialect/Triton/IR/Ops.cpp+++ b/lib/Dialect/Triton/IR/Ops.cpp@@ -1,10 +1,9 @@-#include "triton/Dialect/Triton/IR/Dialect.h"-#include "triton/Dialect/Triton/IR/Types.h"- #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/OperationSupport.h"+#include "triton/Dialect/Triton/IR/Dialect.h"+#include "triton/Dialect/Triton/IR/Types.h" namespace mlir { namespace triton {@@ -38,8 +37,8 @@ static Type getPointerTypeSameShape(Type type) { } // Parser & printer for assembly forms-ParseResult parseLoadOp(OpAsmParser &parser, OperationState &result) {- SmallVector<OpAsmParser::OperandType, 4> allOperands;+ParseResult LoadOp::parse(OpAsmParser &parser, OperationState &result) {+ SmallVector<OpAsmParser::UnresolvedOperand, 4> allOperands; Type resultTypes[1]; SMLoc allOperandLoc = parser.getCurrentLocation(); if (parser.parseOperandList(allOperands) ||@@ -73,18 +72,18 @@ ParseResult parseLoadOp(OpAsmParser &parser, OperationState &result) { return success(); } -void printLoadOp(OpAsmPrinter &printer, LoadOp loadOp) {+void LoadOp::print(OpAsmPrinter &printer) { printer << " ";- printer << loadOp.getOperation()->getOperands();+ printer << getOperation()->getOperands(); // "operand_segment_sizes" can be deduced, so we don't print it.- printer.printOptionalAttrDict(loadOp->getAttrs(),- {loadOp.operand_segment_sizesAttrName()});+ printer.printOptionalAttrDict(getOperation()->getAttrs(),+ {operand_segment_sizesAttrName()}); printer << " : ";- printer.printStrippedAttrOrType(loadOp.result().getType());+ printer.printStrippedAttrOrType(getResult().getType()); } -ParseResult parseStoreOp(OpAsmParser &parser, OperationState &result) {- SmallVector<OpAsmParser::OperandType, 4> allOperands;+ParseResult StoreOp::parse(OpAsmParser &parser, OperationState &result) {+ SmallVector<OpAsmParser::UnresolvedOperand, 4> allOperands; Type valueType; SMLoc allOperandLoc = parser.getCurrentLocation(); if (parser.parseOperandList(allOperands) ||@@ -104,12 +103,12 @@ ParseResult parseStoreOp(OpAsmParser &parser, OperationState &result) { return success(); } -void printStoreOp(OpAsmPrinter &printer, StoreOp storeOp) {+void StoreOp::print(OpAsmPrinter &printer) { printer << " ";- printer << storeOp.getOperation()->getOperands();- printer.printOptionalAttrDict(storeOp->getAttrs(), /*elidedAttrs=*/{});+ printer << getOperation()->getOperands();+ printer.printOptionalAttrDict(getOperation()->getAttrs(), /*elidedAttrs=*/{}); printer << " : ";- printer.printStrippedAttrOrType(storeOp.value().getType());+ printer.printStrippedAttrOrType(value().getType()); } } // namespace triton@@ -319,7 +318,8 @@ OpFoldResult SplatOp::fold(ArrayRef<Attribute> operands) { if (!constOperand) return {}; auto shapedType = getType().cast<ShapedType>();- auto ret = SplatElementsAttr::get(shapedType, {constOperand.getValue()});+ auto ret = SplatElementsAttr::get(+ shapedType, ArrayRef<Attribute>(constOperand.getValue())); return ret; } diff --git a/lib/Dialect/Triton/Transforms/Combine.cpp b/lib/Dialect/Triton/Transforms/Combine.cppindex 2261472170..11570283d6 100644--- a/lib/Dialect/Triton/Transforms/Combine.cpp+++ b/lib/Dialect/Triton/Transforms/Combine.cpp@@ -57,13 +57,13 @@ DenseElementsAttr getConstantValue(Builder &builder, Attribute value, class CombineSelectMaskedLoadPattern : public mlir::RewritePattern { public: CombineSelectMaskedLoadPattern(mlir::MLIRContext *context)- : mlir::RewritePattern(mlir::SelectOp::getOperationName(), 3, context,- {triton::LoadOp::getOperationName()}) {}+ : mlir::RewritePattern(mlir::arith::SelectOp::getOperationName(), 3,+ context, {triton::LoadOp::getOperationName()}) {} mlir::LogicalResult matchAndRewrite(mlir::Operation *op, mlir::PatternRewriter &rewriter) const override {- auto selectOp = llvm::dyn_cast<mlir::SelectOp>(op);+ auto selectOp = llvm::dyn_cast<mlir::arith::SelectOp>(op); if (!selectOp) return mlir::failure(); diff --git a/lib/Dialect/Triton/Transforms/Combine.td b/lib/Dialect/Triton/Transforms/Combine.tdindex 14f286b26e..ded0e346e6 100644--- a/lib/Dialect/Triton/Transforms/Combine.td+++ b/lib/Dialect/Triton/Transforms/Combine.td@@ -1,9 +1,9 @@ #ifndef TRITON_PATTERNS #define TRITON_PATTERNS -include "mlir/Dialect/StandardOps/IR/Ops.td" include "mlir/Dialect/Arithmetic/IR/ArithmeticOps.td" include "triton/Dialect/Triton/IR/TritonOps.td"+include "mlir/IR/PatternBase.td" // AddIOp(DotOp(a, b, c), d) and c==0 => DotOp(a, b, d)diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dialect.cppindex 1fbc609e88..bfc3f3d3da 100644--- a/lib/Dialect/TritonGPU/IR/Dialect.cpp+++ b/lib/Dialect/TritonGPU/IR/Dialect.cpp@@ -1,14 +1,14 @@+#include "triton/Dialect/Triton/IR/Dialect.h"+ #include <numeric> #include "mlir/IR/DialectImplementation.h" #include "mlir/IR/OpImplementation.h" #include "triton/Analysis/Utility.h"-#include "triton/Dialect/Triton/IR/Dialect.h"+#include "triton/Dialect/TritonGPU/IR/Dialect.cpp.inc" #include "triton/Dialect/TritonGPU/IR/Dialect.h" #include "llvm/ADT/TypeSwitch.h" -#include "triton/Dialect/TritonGPU/IR/Dialect.cpp.inc"- using namespace mlir; using namespace mlir::triton::gpu; @@ -366,7 +366,6 @@ template SmallVector<int64_t> SliceEncodingAttr::paddedShape<int64_t>(ArrayRef<int64_t> shape) const; unsigned SliceEncodingAttr::getElemsPerThread(ArrayRef<int64_t> shape) const {- size_t rank = shape.size(); auto parent = getParent(); return ::getElemsPerThread(parent, paddedShape(shape)); }@@ -655,9 +654,9 @@ void DotOperandEncodingAttr::print(mlir::AsmPrinter &printer) const { // InsertSliceAsyncOp //===----------------------------------------------------------------------===// -ParseResult parseInsertSliceAsyncOp(OpAsmParser &parser,- OperationState &result) {- SmallVector<OpAsmParser::OperandType, 8> allOperands;+ParseResult InsertSliceAsyncOp::parse(OpAsmParser &parser,+ OperationState &result) {+ SmallVector<OpAsmParser::UnresolvedOperand, 8> allOperands; Type srcType, dstType; SMLoc allOperandLoc = parser.getCurrentLocation(); if (parser.parseOperandList(allOperands) ||@@ -696,18 +695,16 @@ ParseResult parseInsertSliceAsyncOp(OpAsmParser &parser, return success(); } -void printInsertSliceAsyncOp(OpAsmPrinter &printer,- InsertSliceAsyncOp insertSliceAsyncOp) {+void InsertSliceAsyncOp::print(OpAsmPrinter &printer) { printer << " ";- printer << insertSliceAsyncOp.getOperation()->getOperands();+ printer << getOperation()->getOperands(); // "operand_segment_sizes" can be deduced, so we don't print it.- printer.printOptionalAttrDict(- insertSliceAsyncOp->getAttrs(),- {insertSliceAsyncOp.operand_segment_sizesAttrName()});+ printer.printOptionalAttrDict(getOperation()->getAttrs(),+ {operand_segment_sizesAttrName()}); printer << " : ";- printer.printStrippedAttrOrType(insertSliceAsyncOp.src().getType());+ printer.printStrippedAttrOrType(src().getType()); printer << " -> ";- printer.printStrippedAttrOrType(insertSliceAsyncOp.result().getType());+ printer.printStrippedAttrOrType(result().getType()); } //===----------------------------------------------------------------------===//diff --git a/lib/Dialect/TritonGPU/Transforms/Coalesce.cpp b/lib/Dialect/TritonGPU/Transforms/Coalesce.cppindex 82407980d3..ee6009f44a 100644--- a/lib/Dialect/TritonGPU/Transforms/Coalesce.cpp+++ b/lib/Dialect/TritonGPU/Transforms/Coalesce.cpp@@ -27,7 +27,11 @@ struct CoalescePass : public TritonGPUCoalesceBase<CoalescePass> { auto origType = ptr.getType().cast<RankedTensorType>(); // Get the shape of the tensor. size_t rank = origType.getRank();- AxisInfo info = axisInfo.lookupLatticeElement(ptr)->getValue();+ dataflow::Lattice<AxisInfo> *latticeElement =+ axisInfo.getLatticeElement(ptr);+ AxisInfo info = latticeElement && !latticeElement->isUninitialized()+ ? latticeElement->getValue()+ : AxisInfo(); // Get the contiguity order of `ptr` auto order = argSort(info.getContiguity()); // The desired divisibility is the maximum divisibility@@ -40,7 +44,7 @@ struct CoalescePass : public TritonGPUCoalesceBase<CoalescePass> { for (Value val : op->getResults()) { if (val.getType() != origType) continue;- auto valInfo = axisInfo.lookupLatticeElement(val);+ auto valInfo = axisInfo.getLatticeElement(val); auto currOrder = argSort(valInfo->getValue().getContiguity()); if (order == currOrder) withSameOrder.insert(val);@@ -55,7 +59,7 @@ struct CoalescePass : public TritonGPUCoalesceBase<CoalescePass> { unsigned elemNumBytes = std::max(elemNumBits / 8, 1u); unsigned perThread = 1; for (Value val : withSameOrder) {- AxisInfo info = axisInfo.lookupLatticeElement(val)->getValue();+ AxisInfo info = axisInfo.getLatticeElement(val)->getValue(); unsigned maxMultipleBytes = info.getDivisibility(order[0]); unsigned maxMultiple = std::max(maxMultipleBytes / elemNumBytes, 1u); unsigned maxContig = info.getContiguity(order[0]);@@ -123,8 +127,10 @@ struct CoalescePass : public TritonGPUCoalesceBase<CoalescePass> { void runOnOperation() override { Operation *op = getOperation(); // Run axis info analysis- AxisInfoAnalysis axisInfo(&getContext());- axisInfo.run(op);+ std::unique_ptr<DataFlowSolver> solver = createDataFlowSolver();+ AxisInfoAnalysis *axisInfo = solver->load<AxisInfoAnalysis>();+ if (failed(solver->initializeAndRun(op)))+ return signalPassFailure(); // For each i/o operation, we determine what layout // the pointers should have for best memory coalescing@@ -146,10 +152,10 @@ struct CoalescePass : public TritonGPUCoalesceBase<CoalescePass> { RankedTensorType ty = ptr.getType().template dyn_cast<RankedTensorType>(); if (!ty || !ty.getElementType().isa<PointerType>()) return;- AxisInfo info = axisInfo.lookupLatticeElement(ptr)->getValue();+ AxisInfo info = axisInfo->getLatticeElement(ptr)->getValue(); auto mod = curr->getParentOfType<ModuleOp>(); int numWarps = triton::gpu::TritonGPUDialect::getNumWarps(mod);- auto convertType = getTypeConverter(axisInfo, ptr, numWarps);+ auto convertType = getTypeConverter(*axisInfo, ptr, numWarps); layoutMap[ptr] = convertType; }); diff --git a/lib/Dialect/TritonGPU/Transforms/Combine.cpp b/lib/Dialect/TritonGPU/Transforms/Combine.cppindex efa37ff2dc..089ce3996c 100644--- a/lib/Dialect/TritonGPU/Transforms/Combine.cpp+++ b/lib/Dialect/TritonGPU/Transforms/Combine.cpp@@ -1,6 +1,6 @@ #include "Utility.h" #include "mlir/Analysis/SliceAnalysis.h"-#include "mlir/Dialect/SCF/SCF.h"+#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/BlockAndValueMapping.h" #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/Matchers.h"diff --git a/lib/Dialect/TritonGPU/Transforms/Combine.td b/lib/Dialect/TritonGPU/Transforms/Combine.tdindex 6bf1b14866..6a7b10dbcb 100644--- a/lib/Dialect/TritonGPU/Transforms/Combine.td+++ b/lib/Dialect/TritonGPU/Transforms/Combine.td@@ -3,5 +3,6 @@ include "triton/Dialect/TritonGPU/IR/TritonGPUOps.td" include "triton/Dialect/Triton/IR/TritonOps.td"+include "mlir/IR/PatternBase.td" #endifdiff --git a/lib/Dialect/TritonGPU/Transforms/DecomposeConversions.cpp b/lib/Dialect/TritonGPU/Transforms/DecomposeConversions.cppindex 4bd3bc76bf..b2f8defd81 100644--- a/lib/Dialect/TritonGPU/Transforms/DecomposeConversions.cpp+++ b/lib/Dialect/TritonGPU/Transforms/DecomposeConversions.cpp@@ -1,5 +1,5 @@ #include "mlir/Analysis/SliceAnalysis.h"-#include "mlir/Dialect/SCF/SCF.h"+#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/BlockAndValueMapping.h" #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/Matchers.h"diff --git a/lib/Dialect/TritonGPU/Transforms/Pipeline.cpp b/lib/Dialect/TritonGPU/Transforms/Pipeline.cppindex 9b2f42231e..85f746c1dc 100644--- a/lib/Dialect/TritonGPU/Transforms/Pipeline.cpp+++ b/lib/Dialect/TritonGPU/Transforms/Pipeline.cpp@@ -2,6 +2,7 @@ #include "mlir/IR/BlockAndValueMapping.h" #include "mlir/IR/TypeUtilities.h" #include "triton/Analysis/AxisInfo.h"+#include "triton/Analysis/Utility.h" #include "triton/Dialect/TritonGPU/IR/Dialect.h" #include "triton/Dialect/TritonGPU/Transforms/Passes.h" @@ -160,15 +161,18 @@ ttg::AllocTensorOp LoopPipeliner::allocateEmptyBuffer(Operation *op, LogicalResult LoopPipeliner::initialize() { Block *loop = forOp.getBody(); - AxisInfoAnalysis axisInfoAnalysis(forOp.getContext());- axisInfoAnalysis.run(forOp->getParentOfType<ModuleOp>());+ std::unique_ptr<DataFlowSolver> solver = createDataFlowSolver();+ AxisInfoAnalysis *axisInfoAnalysis = solver->load<AxisInfoAnalysis>();+ if (failed(solver->initializeAndRun(forOp->getParentOfType<ModuleOp>()))) {+ return failure();+ } // can we use forOp.walk(...) here? SmallVector<triton::LoadOp, 2> allLoads; for (Operation &op : *loop) if (auto loadOp = dyn_cast<triton::LoadOp>(&op)) { auto ptr = loadOp.ptr();- unsigned vec = axisInfoAnalysis.getPtrContiguity(ptr);+ unsigned vec = axisInfoAnalysis->getPtrContiguity(ptr); auto tensorTy = ptr.getType().dyn_cast<RankedTensorType>(); if (!tensorTy) continue;diff --git a/lib/Dialect/TritonGPU/Transforms/ReorderInstructions.cpp b/lib/Dialect/TritonGPU/Transforms/ReorderInstructions.cppindex 0e7dbe5264..b95a4f50a6 100644--- a/lib/Dialect/TritonGPU/Transforms/ReorderInstructions.cpp+++ b/lib/Dialect/TritonGPU/Transforms/ReorderInstructions.cpp@@ -1,5 +1,5 @@ #include "mlir/Analysis/SliceAnalysis.h"-#include "mlir/Dialect/SCF/SCF.h"+#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/BlockAndValueMapping.h" #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/Matchers.h"diff --git a/lib/Dialect/TritonGPU/Transforms/TritonGPUConversion.cpp b/lib/Dialect/TritonGPU/Transforms/TritonGPUConversion.cppindex 37ac710995..762e887f36 100644--- a/lib/Dialect/TritonGPU/Transforms/TritonGPUConversion.cpp+++ b/lib/Dialect/TritonGPU/Transforms/TritonGPUConversion.cpp@@ -82,12 +82,12 @@ TritonGPUConversionTarget::TritonGPUConversionTarget( scf::ReduceReturnOp>(); addDynamicallyLegalDialect<arith::ArithmeticDialect, math::MathDialect,- triton::TritonDialect, StandardOpsDialect,- scf::SCFDialect>([&](Operation *op) {- if (typeConverter.isLegal(op))- return true;- return false;- });+ triton::TritonDialect, scf::SCFDialect>(+ [&](Operation *op) {+ if (typeConverter.isLegal(op))+ return true;+ return false;+ }); // We have requirements for the data layouts addDynamicallyLegalOp<triton::DotOp>([](triton::DotOp dotOp) -> bool {diff --git a/lib/Dialect/TritonGPU/Transforms/UpdateMmaForVolta.cpp b/lib/Dialect/TritonGPU/Transforms/UpdateMmaForVolta.cppindex c229104286..c911fd4a5c 100644--- a/lib/Dialect/TritonGPU/Transforms/UpdateMmaForVolta.cpp+++ b/lib/Dialect/TritonGPU/Transforms/UpdateMmaForVolta.cpp@@ -1,5 +1,5 @@ #include "Utility.h"-#include "mlir/Dialect/SCF/SCF.h"+#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/Matchers.h" #include "mlir/IR/PatternMatch.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h"@@ -118,8 +118,8 @@ void setOpResultType(Operation *op, ArrayRef<Type> newTypes) { .get("value") .dyn_cast<mlir::DenseElementsAttr>(); if (attr) {- auto newAttr = mlir::DenseElementsAttr::getFromRawBuffer(- newType, attr.getRawData(), true);+ auto newAttr =+ mlir::DenseElementsAttr::getFromRawBuffer(newType, attr.getRawData()); op->setAttr("value", newAttr); } }diff --git a/lib/Dialect/TritonGPU/Transforms/Utility.cpp b/lib/Dialect/TritonGPU/Transforms/Utility.cppindex ed15f02f67..6400f1633a 100644--- a/lib/Dialect/TritonGPU/Transforms/Utility.cpp+++ b/lib/Dialect/TritonGPU/Transforms/Utility.cpp@@ -1,5 +1,5 @@ #include "Utility.h"-#include "mlir/Dialect/SCF/SCF.h"+#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/BlockAndValueMapping.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h" diff --git a/lib/Target/LLVMIR/CMakeLists.txt b/lib/Target/LLVMIR/CMakeLists.txtindex f1bbd0bf4e..ac8973ad19 100644--- a/lib/Target/LLVMIR/CMakeLists.txt+++ b/lib/Target/LLVMIR/CMakeLists.txt@@ -6,8 +6,7 @@ add_mlir_translation_library(TritonLLVMIR LINK_LIBS PUBLIC MLIRIR- MLIRLLVMIR- MLIRSCFToStandard+ MLIRLLVMDialect MLIRSupport MLIRTargetLLVMIRExport )diff --git a/lib/Target/PTX/PTXTranslation.cpp b/lib/Target/PTX/PTXTranslation.cppindex 4cb0d8193c..6a5453a6e7 100644--- a/lib/Target/PTX/PTXTranslation.cpp+++ b/lib/Target/PTX/PTXTranslation.cpp@@ -1,11 +1,14 @@ #include "triton/Target/PTX/PTXTranslation.h" #include "triton/Target/LLVMIR/LLVMIRTranslation.h"+#include <optional> #include "llvm/IR/IRBuilder.h" #include "llvm/IR/LegacyPassManager.h" #include "llvm/IR/Module.h" #include "llvm/IR/Verifier.h" #include "llvm/MC/TargetRegistry.h"+#include "llvm/Pass.h"+#include "llvm/Support/CommandLine.h" #include "llvm/Support/TargetSelect.h" #include "llvm/Target/TargetMachine.h" diff --git a/python/setup.py b/python/setup.pyindex 2ac3accd25..4530b36714 100644--- a/python/setup.py+++ b/python/setup.py@@ -57,19 +57,10 @@ def get_pybind11_package_info(): def get_llvm_package_info(): # download if nothing is installed system = platform.system()- if system == "Darwin":- system_suffix = "apple-darwin"- elif system == "Linux":- vglibc = tuple(map(int, platform.libc_ver()[1].split('.')))- vglibc = vglibc[0] * 100 + vglibc[1]- linux_suffix = 'ubuntu-18.04' if vglibc > 217 else 'centos-7'- system_suffix = f"linux-gnu-{linux_suffix}"- else:- raise RuntimeError(f"unsupported system: {system}")+ system_suffix = {"Linux": "linux-gnu-ubuntu-18.04", "Darwin": "apple-darwin"}[system] use_assert_enabled_llvm = check_env_flag("TRITON_USE_ASSERT_ENABLED_LLVM", "False")- release_suffix = "assert" if use_assert_enabled_llvm else "release"- name = f'llvm+mlir-14.0.6-x86_64-{system_suffix}-{release_suffix}'- url = f"https://github.com/ptillet/triton-llvm-releases/releases/download/llvm-14.0.6-f28c006a5895/{name}.tar.xz"+ name = 'llvm+mlir-15.0.7-x86_64-{}-{}'.format(system_suffix, "assert" if use_assert_enabled_llvm else "release")+ url = "https://github.com/ptillet/triton-llvm-releases/releases/download/llvm-15.0.7-8dfdcc7b7bf6/{}.tar.xz".format(name) return Package("llvm", name, url, "lib", "LLVM_INCLUDE_DIRS", "LLVM_LIBRARY_DIR", "LLVM_SYSPATH") diff --git a/python/src/triton.cc b/python/src/triton.ccindex c40b117a55..f190eacc34 100644--- a/python/src/triton.cc+++ b/python/src/triton.cc@@ -8,9 +8,10 @@ #include "mlir/Pass/PassManager.h" #include "mlir/Transforms/Passes.h" -#include "mlir/Parser.h"+#include "mlir/Parser/Parser.h" #include "mlir/Support/FileUtilities.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "triton/Analysis/Allocation.h" #include "triton/Conversion/TritonGPUToLLVM/TritonGPUToLLVMPass.h"@@ -195,7 +196,7 @@ void init_triton_ir(py::module &&m) { std::string attrName = name + "_arg" + std::to_string(id); mlir::Block *owner = arg.getOwner(); if (owner->isEntryBlock() &&- !mlir::isa<mlir::FuncOp>(owner->getParentOp())) {+ !mlir::isa<mlir::func::FuncOp>(owner->getParentOp())) { owner->getParentOp()->setAttr(attrName, attr); } }@@ -348,7 +349,7 @@ void init_triton_ir(py::module &&m) { return str; }) .def("push_back",- [](mlir::ModuleOp &self, mlir::FuncOp &funcOp) -> void {+ [](mlir::ModuleOp &self, mlir::func::FuncOp &funcOp) -> void { self.push_back(funcOp); }) .def("has_function",@@ -358,16 +359,18 @@ void init_triton_ir(py::module &&m) { return false; }) .def("get_function",- [](mlir::ModuleOp &self, std::string &funcName) -> mlir::FuncOp {- return self.lookupSymbol<mlir::FuncOp>(funcName);- })- .def("get_single_function", [](mlir::ModuleOp &self) -> mlir::FuncOp {- llvm::SmallVector<mlir::FuncOp> funcs;- self.walk([&](mlir::FuncOp func) { funcs.push_back(func); });- if (funcs.size() != 1)- throw std::runtime_error("Expected a single function");- return funcs[0];- });+ [](mlir::ModuleOp &self,+ std::string &funcName) -> mlir::func::FuncOp {+ return self.lookupSymbol<mlir::func::FuncOp>(funcName);+ })+ .def("get_single_function",+ [](mlir::ModuleOp &self) -> mlir::func::FuncOp {+ llvm::SmallVector<mlir::func::FuncOp> funcs;+ self.walk([&](mlir::func::FuncOp func) { funcs.push_back(func); });+ if (funcs.size() != 1)+ throw std::runtime_error("Expected a single function");+ return funcs[0];+ }); m.def("make_attr", [](const std::vector<int> &values, mlir::MLIRContext &context) {@@ -388,47 +391,48 @@ void init_triton_ir(py::module &&m) { registry.insert<mlir::triton::TritonDialect, mlir::triton::gpu::TritonGPUDialect, mlir::math::MathDialect, mlir::arith::ArithmeticDialect,- mlir::StandardOpsDialect, mlir::scf::SCFDialect>();+ mlir::func::FuncDialect, mlir::scf::SCFDialect>(); context.appendDialectRegistry(registry); context.loadAllAvailableDialects(); // parse module- mlir::OwningOpRef<mlir::ModuleOp> module(- mlir::parseSourceFile(inputFilename, &context));+ mlir::OwningOpRef<mlir::ModuleOp> module =+ mlir::parseSourceFile<mlir::ModuleOp>(inputFilename, &context);+ if (!module)+ throw std::runtime_error("Parse MLIR file failed."); // locations are incompatible with ptx < 7.5 ! module->walk([](mlir::Operation *op) { op->setLoc(mlir::UnknownLoc::get(op->getContext())); });- if (!module)- throw std::runtime_error("Parse MLIR file failed."); return module->clone(); }, ret::take_ownership); - py::class_<mlir::FuncOp, mlir::OpState>(m, "function")+ py::class_<mlir::func::FuncOp, mlir::OpState>(m, "function") // .def_property_readonly("attrs", &ir::function::attrs) // .def("add_attr", &ir::function::add_attr); .def("args",- [](mlir::FuncOp &self, unsigned idx) -> mlir::BlockArgument {+ [](mlir::func::FuncOp &self, unsigned idx) -> mlir::BlockArgument { return self.getArgument(idx); }) .def( "add_entry_block",- [](mlir::FuncOp &self) -> mlir::Block * {+ [](mlir::func::FuncOp &self) -> mlir::Block * { return self.addEntryBlock(); }, ret::reference) .def( "set_arg_attr",- [](mlir::FuncOp &self, int arg_no, const std::string &name, int val) {+ [](mlir::func::FuncOp &self, int arg_no, const std::string &name,+ int val) { // set arg attributes "name" to value "val" auto attrTy = mlir::IntegerType::get(self.getContext(), 32); self.setArgAttr(arg_no, name, mlir::IntegerAttr::get(attrTy, val)); }, ret::reference)- .def_property_readonly("type", &mlir::FuncOp::getType)- .def("reset_type", &mlir::FuncOp::setType);+ .def_property_readonly("type", &mlir::func::FuncOp::getFunctionType)+ .def("reset_type", &mlir::func::FuncOp::setType); py::class_<mlir::OpBuilder::InsertPoint>(m, "InsertPoint"); @@ -445,13 +449,13 @@ void init_triton_ir(py::module &&m) { .def("ret", [](mlir::OpBuilder &self, std::vector<mlir::Value> &vals) -> void { auto loc = self.getUnknownLoc();- self.create<mlir::ReturnOp>(loc, vals);+ self.create<mlir::func::ReturnOp>(loc, vals); }) .def("call",- [](mlir::OpBuilder &self, mlir::FuncOp &func,+ [](mlir::OpBuilder &self, mlir::func::FuncOp &func, std::vector<mlir::Value> &args) -> mlir::OpState { auto loc = self.getUnknownLoc();- return self.create<mlir::CallOp>(loc, func, args);+ return self.create<mlir::func::CallOp>(loc, func, args); }) // insertion block/point .def("set_insertion_point_to_start",@@ -618,15 +622,16 @@ void init_triton_ir(py::module &&m) { .def("get_or_insert_function", [](mlir::OpBuilder &self, mlir::ModuleOp &module, std::string &funcName, mlir::Type &funcType,- std::string &visibility) -> mlir::FuncOp {+ std::string &visibility) -> mlir::func::FuncOp { if (mlir::Operation *funcOperation = module.lookupSymbol(funcName))- return llvm::dyn_cast<mlir::FuncOp>(funcOperation);+ return llvm::dyn_cast<mlir::func::FuncOp>(funcOperation); auto loc = self.getUnknownLoc(); if (auto funcTy = funcType.dyn_cast<mlir::FunctionType>()) { llvm::SmallVector<mlir::NamedAttribute> attrs = { mlir::NamedAttribute(self.getStringAttr("sym_visibility"), self.getStringAttr(visibility))};- return self.create<mlir::FuncOp>(loc, funcName, funcTy, attrs);+ return self.create<mlir::func::FuncOp>(loc, funcName, funcTy,+ attrs); } throw std::runtime_error("invalid function type"); })@@ -658,15 +663,15 @@ void init_triton_ir(py::module &&m) { [](mlir::OpBuilder &self, mlir::Value condition, mlir::Block *trueDest, mlir::Block *falseDest) { auto loc = self.getUnknownLoc();- self.create<mlir::CondBranchOp>(loc, condition, trueDest,- falseDest);+ self.create<mlir::cf::CondBranchOp>(loc, condition, trueDest,+ falseDest); return; }) .def("create_branch", [](mlir::OpBuilder &self, mlir::Block *dest, std::vector<mlir::Value> &args) { auto loc = self.getUnknownLoc();- self.create<mlir::BranchOp>(loc, dest, args);+ self.create<mlir::cf::BranchOp>(loc, dest, args); return; }) // Structured control flow@@ -792,14 +797,14 @@ void init_triton_ir(py::module &&m) { .def("create_to_index", [](mlir::OpBuilder &self, mlir::Value &input) -> mlir::Value { auto loc = self.getUnknownLoc();- return self.create<mlir::arith::IndexCastOp>(loc, input,- self.getIndexType());+ return self.create<mlir::arith::IndexCastOp>(+ loc, self.getIndexType(), input); }) .def("create_index_to_si", [](mlir::OpBuilder &self, mlir::Value &input) -> mlir::Value { auto loc = self.getUnknownLoc();- return self.create<mlir::arith::IndexCastOp>(loc, input,- self.getI32Type());+ return self.create<mlir::arith::IndexCastOp>(+ loc, self.getI32Type(), input); }) .def("create_fmul", [](mlir::OpBuilder &self, mlir::Value &lhs,@@ -1316,8 +1321,8 @@ void init_triton_ir(py::module &&m) { [](mlir::OpBuilder &self, mlir::Value &condition, mlir::Value &trueValue, mlir::Value &falseValue) -> mlir::Value { auto loc = self.getUnknownLoc();- return self.create<mlir::SelectOp>(loc, condition, trueValue,- falseValue);+ return self.create<mlir::arith::SelectOp>(loc, condition,+ trueValue, falseValue); }) .def("create_printf", [](mlir::OpBuilder &self, const std::string &prefix,@@ -1429,7 +1434,7 @@ void init_triton_ir(py::module &&m) { self.addPass(mlir::triton::createConvertTritonGPUToLLVMPass()); }) .def("add_scf_to_cfg", [](mlir::PassManager &self) {- self.addPass(mlir::createLowerToCFGPass());+ self.addPass(mlir::createConvertSCFToCFPass()); }); } diff --git a/python/test/unit/language/test_core.py b/python/test/unit/language/test_core.pyindex 432544a8a4..018f544714 100644--- a/python/test/unit/language/test_core.py+++ b/python/test/unit/language/test_core.py@@ -1918,7 +1918,7 @@ def test_convert2d(dtype, shape, src_layout, dst_layout, device='cuda'): #dst = {dst_layout} """ + """ module attributes {"triton_gpu.num-warps" = 4 : i32} {- func public @kernel_0d1d(%arg0: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f16> {tt.divisibility = 16 : i32}) {+ func.func public @kernel_0d1d(%arg0: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f16> {tt.divisibility = 16 : i32}) { %cst = arith.constant dense<128> : tensor<128x1xi32, #src> %0 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32, #triton_gpu.slice<{dim = 1, parent = #src}>> %1 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32, #triton_gpu.slice<{dim = 0, parent = #src}>>diff --git a/python/triton/compiler.py b/python/triton/compiler.pyindex 5d167634df..c36589037c 100644--- a/python/triton/compiler.py+++ b/python/triton/compiler.py@@ -1514,14 +1514,14 @@ def make_hash(fn, **kwargs): return hashlib.md5((Path(fn).read_text() + triton.runtime.jit.version_key()).encode("utf-8")).hexdigest() -# - ^\s*func\s+ : match the start of the string, any leading whitespace, the keyword func,+# - ^\s*func\.func\s+ : match the start of the string, any leading whitespace, the keyword func, # and any following whitespace # - (public\s+)? : optionally match the keyword public and any following whitespace # - (@\w+) : match an @ symbol followed by one or more word characters # (letters, digits, or underscores), and capture it as group 1 (the function name) # - (\((?:%\w+: \S+(?: \{\S+ = \S+ : \S+\})?(?:, )?)*\)) : match a pair of parentheses enclosing # zero or more arguments separated by commas, and capture it as group 2 (the argument list)-mlir_prototype_pattern = r'^\s*func\s+(?:public\s+)?(@\w+)(\((?:%\w+: \S+(?: \{\S+ = \S+ : \S+\})?(?:, )?)*\))\s*\{\s*$'+mlir_prototype_pattern = r'^\s*func\.func\s+(?:public\s+)?(@\w+)(\((?:%\w+: \S+(?: \{\S+ = \S+ : \S+\})?(?:, )?)*\))\s*\{\s*$' ptx_prototype_pattern = r"\.(?:visible|extern)\s+\.(?:entry|func)\s+(\w+)\s*\(([^)]*)\)" prototype_pattern = { "ttir": mlir_prototype_pattern,diff --git a/test/Analysis/test-alias.mlir b/test/Analysis/test-alias.mlirindex b3d5673f85..bb21615e68 100644--- a/test/Analysis/test-alias.mlir+++ b/test/Analysis/test-alias.mlir@@ -11,7 +11,7 @@ // CHECK-LABEL: matmul_loop // There shouldn't be any aliasing with the dot op encoding.-func @matmul_loop(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) {+func.func @matmul_loop(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) { %a_ptr_init = tt.broadcast %A : (!tt.ptr<f16>) -> tensor<128x32x!tt.ptr<f16>, #AL> %b_ptr_init = tt.broadcast %B : (!tt.ptr<f16>) -> tensor<32x128x!tt.ptr<f16>, #BL> %a_mask = arith.constant dense<true> : tensor<128x32xi1, #AL>@@ -36,7 +36,7 @@ func @matmul_loop(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B } // CHECK-LABEL: alloc-func @alloc(%A : !tt.ptr<f16>) {+func.func @alloc(%A : !tt.ptr<f16>) { // CHECK: %cst -> %cst %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> %cst1 = arith.constant dense<0.000000e+00> : tensor<16x32xf16, #AL>@@ -46,7 +46,7 @@ func @alloc(%A : !tt.ptr<f16>) { } // CHECK-LABEL: convert-func @convert(%A : !tt.ptr<f16>) {+func.func @convert(%A : !tt.ptr<f16>) { %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #AL> // CHECK: %0 -> %0 %cst1 = triton_gpu.convert_layout %cst0 : (tensor<16x16xf16, #AL>) -> tensor<16x16xf16, #A_SHARED>@@ -54,7 +54,7 @@ func @convert(%A : !tt.ptr<f16>) { } // CHECK-LABEL: trans-func @trans(%A : !tt.ptr<f16>) {+func.func @trans(%A : !tt.ptr<f16>) { // CHECK: %cst -> %cst %tensor = arith.constant dense<0.000000e+00> : tensor<16x32xf16, #A_SHARED> // CHECK: %0 -> %cst@@ -63,7 +63,7 @@ func @trans(%A : !tt.ptr<f16>) { } // CHECK-LABEL: insert_slice_async-func @insert_slice_async(%A : !tt.ptr<f16>, %i1 : i1) {+func.func @insert_slice_async(%A : !tt.ptr<f16>, %i1 : i1) { %a_ptr = tt.broadcast %A : (!tt.ptr<f16>) -> tensor<16x16x!tt.ptr<f16>, #AL> %mask = tt.splat %i1 : (i1) -> tensor<16x16xi1, #AL> %other = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #AL>@@ -76,7 +76,7 @@ func @insert_slice_async(%A : !tt.ptr<f16>, %i1 : i1) { } // CHECK-LABEL: insert_slice-func @insert_slice(%A : !tt.ptr<f16>, %i1 : i1) {+func.func @insert_slice(%A : !tt.ptr<f16>, %i1 : i1) { %a_ptr = tt.broadcast %A : (!tt.ptr<f16>) -> tensor<16x16x!tt.ptr<f16>, #AL> %mask = tt.splat %i1 : (i1) -> tensor<16x16xi1, #AL> %other = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #AL>@@ -90,7 +90,7 @@ func @insert_slice(%A : !tt.ptr<f16>, %i1 : i1) { } // CHECK-LABEL: extract_slice-func @extract_slice(%A : !tt.ptr<f16>) {+func.func @extract_slice(%A : !tt.ptr<f16>) { // CHECK: %cst -> %cst %cst0 = arith.constant dense<0.000000e+00> : tensor<1x16x16xf16, #A_SHARED> %index = arith.constant 0 : index@@ -100,7 +100,7 @@ func @extract_slice(%A : !tt.ptr<f16>) { } // CHECK-LABEL: if_cat-func @if_cat(%i1 : i1) {+func.func @if_cat(%i1 : i1) { // CHECK: %cst -> %cst %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> // CHECK: %cst_0 -> %cst_0@@ -119,7 +119,7 @@ func @if_cat(%i1 : i1) { } // CHECK-LABEL: if_alias-func @if_alias(%i1 : i1) {+func.func @if_alias(%i1 : i1) { // CHECK: %cst -> %cst %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> // CHECK-NEXT: %cst_0 -> %cst_0@@ -134,7 +134,7 @@ func @if_alias(%i1 : i1) { } // CHECK-LABEL: for-func @for(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) {+func.func @for(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) { // CHECK: %cst -> %cst %a_shared_init = arith.constant dense<0.00e+00> : tensor<128x32xf16, #A_SHARED> // CHECK-NEXT: %cst_0 -> %cst_0@@ -154,7 +154,7 @@ func @for(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.p } // CHECK-LABEL: for_if-func @for_if(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>, %i1 : i1) {+func.func @for_if(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>, %i1 : i1) { // CHECK: %cst -> %cst %a_shared_init = arith.constant dense<0.00e+00> : tensor<128x32xf16, #A_SHARED> // CHECK-NEXT: %cst_0 -> %cst_0@@ -180,7 +180,7 @@ func @for_if(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !t } // CHECK-LABEL: for_if_for-func @for_if_for(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>, %i1 : i1) {+func.func @for_if_for(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>, %i1 : i1) { // CHECK: %cst -> %cst %a_shared_init = arith.constant dense<0.00e+00> : tensor<128x32xf16, #A_SHARED> // CHECK-NEXT: %cst_0 -> %cst_0diff --git a/test/Analysis/test-alignment.mlir b/test/Analysis/test-alignment.mlirindex 0ab34c7a78..af8ea6f856 100644--- a/test/Analysis/test-alignment.mlir+++ b/test/Analysis/test-alignment.mlir@@ -1,288 +1,288 @@-// RUN: triton-opt %s -test-print-alignment -split-input-file 2>&1 | FileCheck %s+// RUN: triton-opt %s -test-print-alignment -split-input-file -o %t 2>&1 | FileCheck %s -// CHECK-LABEL: cast-func @cast() {- // CHECK: Contiguity: [1] ; Divisibility: [1] ; Constancy: [1] ; ConstantValue: [1]+// CHECK-LABEL: @cast+func.func @cast() {+ // CHECK: contiguity = [1], divisibility = [1], constancy = [1], constant_value = 1 %cst = arith.constant 1 : i32- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [1] ; ConstantValue: [1]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [1], constant_value = 1 %0 = arith.extsi %cst : i32 to i64- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [128] ; ConstantValue: [1]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [128], constant_value = 1 %cst_tensor = arith.constant dense<1> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [128] ; ConstantValue: [1]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [128], constant_value = 1 %1 = tt.bitcast %cst_tensor : tensor<128xi32> -> tensor<128xi64> return } // ----- -// CHECK-LABEL: add-func @add() {- // CHECK: Contiguity: [128] ; Divisibility: [1073741824] ; Constancy: [1] ; ConstantValue: [None]+// CHECK-LABEL: @add+func.func @add() {+ // CHECK: contiguity = [128], divisibility = [1073741824], constancy = [1], constant_value = <none> %0 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [128] ; ConstantValue: [1]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [128], constant_value = 1 %1 = arith.constant dense<1> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [128] ; Divisibility: [1] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [128], divisibility = [1], constancy = [1], constant_value = <none> %2 = arith.addi %0, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [128] ; ConstantValue: [127]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [128], constant_value = 127 %3 = arith.constant dense<127> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [128] ; Constancy: [128] ; ConstantValue: [128]+ // CHECK-NEXT: contiguity = [1], divisibility = [128], constancy = [128], constant_value = 128 %4 = arith.addi %1, %3 : tensor<128xi32> return } // ----- -// CHECK-LABEL: sub-func @sub() {- // CHECK: Contiguity: [128] ; Divisibility: [1073741824] ; Constancy: [1] ; ConstantValue: [None]+// CHECK-LABEL: @sub+func.func @sub() {+ // CHECK: contiguity = [128], divisibility = [1073741824], constancy = [1], constant_value = <none> %0 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [128] ; ConstantValue: [1]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [128], constant_value = 1 %1 = arith.constant dense<1> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [128] ; Divisibility: [1] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [128], divisibility = [1], constancy = [1], constant_value = <none> %2 = arith.subi %0, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [128] ; ConstantValue: [129]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [128], constant_value = 129 %3 = arith.constant dense<129> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [128] ; Constancy: [128] ; ConstantValue: [128]+ // CHECK-NEXT: contiguity = [1], divisibility = [128], constancy = [128], constant_value = 128 %4 = arith.subi %3, %1 : tensor<128xi32> return } // ----- -// CHECK-LABEL: mul-func @mul() {- // CHECK: Contiguity: [128] ; Divisibility: [1073741824] ; Constancy: [1] ; ConstantValue: [None]+// CHECK-LABEL: @mul+func.func @mul() {+ // CHECK: contiguity = [128], divisibility = [1073741824], constancy = [1], constant_value = <none> %0 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [128] ; ConstantValue: [1]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [128], constant_value = 1 %1 = arith.constant dense<1> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [128] ; Divisibility: [1073741824] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [128], divisibility = [1073741824], constancy = [1], constant_value = <none> %2 = arith.muli %0, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [128] ; Constancy: [128] ; ConstantValue: [128]+ // CHECK-NEXT: contiguity = [1], divisibility = [128], constancy = [128], constant_value = 128 %3 = arith.constant dense<128> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [128] ; Constancy: [128] ; ConstantValue: [128]+ // CHECK-NEXT: contiguity = [1], divisibility = [128], constancy = [128], constant_value = 128 %4 = arith.muli %3, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [2] ; Constancy: [128] ; ConstantValue: [2]+ // CHECK-NEXT: contiguity = [1], divisibility = [2], constancy = [128], constant_value = 2 %5 = arith.constant dense<2> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [256] ; Constancy: [128] ; ConstantValue: [256]+ // CHECK-NEXT: contiguity = [1], divisibility = [256], constancy = [128], constant_value = 256 %6 = arith.muli %4, %5 : tensor<128xi32> return } // ----- -// CHECK-LABEL: div-func @div() {- // CHECK: Contiguity: [128] ; Divisibility: [1073741824] ; Constancy: [1] ; ConstantValue: [None]+// CHECK-LABEL: @div+func.func @div() {+ // CHECK: contiguity = [128], divisibility = [1073741824], constancy = [1], constant_value = <none> %0 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [128] ; ConstantValue: [1]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [128], constant_value = 1 %1 = arith.constant dense<1> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [128] ; Divisibility: [1073741824] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [128], divisibility = [1073741824], constancy = [1], constant_value = <none> %2 = arith.divsi %0, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [1], constant_value = <none> %3 = arith.divui %1, %0 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [64] ; Constancy: [128] ; ConstantValue: [64]+ // CHECK-NEXT: contiguity = [1], divisibility = [64], constancy = [128], constant_value = 64 %4 = arith.constant dense<64> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [16777216] ; Constancy: [64] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [16777216], constancy = [64], constant_value = <none> %5 = arith.divsi %0, %4 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [1], constant_value = <none> %6 = arith.divsi %4, %0 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [64] ; Constancy: [128] ; ConstantValue: [64]+ // CHECK-NEXT: contiguity = [1], divisibility = [64], constancy = [128], constant_value = 64 %7 = arith.divsi %4, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [2] ; Constancy: [128] ; ConstantValue: [66]+ // CHECK-NEXT: contiguity = [1], divisibility = [2], constancy = [128], constant_value = 66 %8 = arith.constant dense<66> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [2] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [2], constant_value = <none> %9 = arith.divui %0, %8 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [128] ; Divisibility: [8192] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [128], divisibility = [8192], constancy = [1], constant_value = <none> %10 = tt.make_range {end = 8320 : i32, start = 8192 : i32} : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [128] ; Constancy: [64] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [128], constancy = [64], constant_value = <none> %11 = arith.divsi %10, %4 : tensor<128xi32>- return + return } // ----- -// CHECK-LABEL: rem-func @rem() {- // CHECK: Contiguity: [128] ; Divisibility: [1073741824] ; Constancy: [1] ; ConstantValue: [None]+// CHECK-LABEL: @rem+func.func @rem() {+ // CHECK: contiguity = [128], divisibility = [1073741824], constancy = [1], constant_value = <none> %0 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [128] ; ConstantValue: [1]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [128], constant_value = 1 %1 = arith.constant dense<1> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [4611686018427387904] ; Constancy: [128] ; ConstantValue: [0]+ // CHECK-NEXT: contiguity = [1], divisibility = [4611686018427387904], constancy = [128], constant_value = 0 %2 = arith.remsi %0, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [1], constant_value = <none> %3 = arith.remui %1, %0 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [64] ; Constancy: [128] ; ConstantValue: [64]+ // CHECK-NEXT: contiguity = [1], divisibility = [64], constancy = [128], constant_value = 64 %4 = arith.constant dense<64> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [64] ; Divisibility: [64] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [64], divisibility = [64], constancy = [1], constant_value = <none> %5 = arith.remsi %0, %4 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [64] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [64], constancy = [1], constant_value = <none> %6 = arith.remsi %4, %0 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [2] ; Constancy: [128] ; ConstantValue: [66]+ // CHECK-NEXT: contiguity = [1], divisibility = [2], constancy = [128], constant_value = 66 %7 = arith.constant dense<66> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [2] ; Divisibility: [2] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [2], divisibility = [2], constancy = [1], constant_value = <none> %8 = arith.remui %0, %7 : tensor<128xi32>- return + return } // ----- -// CHECK-LABEL: broadcast-func @broadcast() {- // CHECK: Contiguity: [1] ; Divisibility: [64] ; Constancy: [128] ; ConstantValue: [64]+// CHECK-LABEL: @broadcast+func.func @broadcast() {+ // CHECK: contiguity = [1], divisibility = [64], constancy = [128], constant_value = 64 %0 = arith.constant dense<64> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1, 1] ; Divisibility: [64, 1] ; Constancy: [128, 1] ; ConstantValue: [64]+ // CHECK-NEXT: contiguity = [1, 1], divisibility = [64, 1], constancy = [128, 1], constant_value = 64 %1 = tt.expand_dims %0 {axis = 1 : i32} : (tensor<128xi32>) -> tensor<128x1xi32>- // CHECK-NEXT: Contiguity: [1, 1] ; Divisibility: [64, 1] ; Constancy: [128, 128] ; ConstantValue: [64]+ // CHECK-NEXT: contiguity = [1, 1], divisibility = [64, 1], constancy = [128, 128], constant_value = 64 %2 = tt.broadcast %1 : (tensor<128x1xi32>) -> tensor<128x128xi32> return } // ----- -// CHECK-LABEL: splat-func @splat(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}) {- // CHECK: Contiguity: [1, 1] ; Divisibility: [16, 16] ; Constancy: [128, 128] ; ConstantValue: [None]+// CHECK-LABEL: @splat+func.func @splat(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}) {+ // CHECK: contiguity = [1, 1], divisibility = [16, 16], constancy = [128, 128], constant_value = <none> %0 = tt.splat %arg0 : (!tt.ptr<f32>) -> tensor<128x128x!tt.ptr<f32>> return } // ----- -// CHECK-LABEL: cmp-func @cmp() {- // CHECK: Contiguity: [128] ; Divisibility: [1073741824] ; Constancy: [1] ; ConstantValue: [None]+// CHECK-LABEL: @cmp+func.func @cmp() {+ // CHECK: contiguity = [128], divisibility = [1073741824], constancy = [1], constant_value = <none> %0 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [4611686018427387904] ; Constancy: [128] ; ConstantValue: [0]+ // CHECK-NEXT: contiguity = [1], divisibility = [4611686018427387904], constancy = [128], constant_value = 0 %1 = arith.constant dense<0> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [128] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [128], constant_value = <none> %2 = arith.cmpi eq, %0, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [128] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [128], constant_value = <none> %3 = arith.cmpi slt, %0, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [1], constant_value = <none> %4 = arith.cmpi sle, %0, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [128] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [128], constant_value = <none> %5 = arith.cmpi sge, %0, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [8] ; Constancy: [128] ; ConstantValue: [8]+ // CHECK-NEXT: contiguity = [1], divisibility = [8], constancy = [128], constant_value = 8 %6 = arith.constant dense<8> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [8] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [8], constant_value = <none> %7 = arith.cmpi sgt, %0, %6 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [128] ; ConstantValue: [0]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [128], constant_value = 0 %8 = arith.cmpi sgt, %1, %6 : tensor<128xi32> return } // ----- -// CHECK-LABEL: logic-func @logic() {- // CHECK: Contiguity: [128] ; Divisibility: [1073741824] ; Constancy: [1] ; ConstantValue: [None]+// CHECK-LABEL: @logic+func.func @logic() {+ // CHECK: contiguity = [128], divisibility = [1073741824], constancy = [1], constant_value = <none> %0 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [64] ; Constancy: [128] ; ConstantValue: [64]+ // CHECK-NEXT: contiguity = [1], divisibility = [64], constancy = [128], constant_value = 64 %1 = arith.constant dense<64> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [16777216] ; Constancy: [64] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [16777216], constancy = [64], constant_value = <none> %2 = arith.divsi %0, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [8] ; Constancy: [128] ; ConstantValue: [8]+ // CHECK-NEXT: contiguity = [1], divisibility = [8], constancy = [128], constant_value = 8 %3 = arith.constant dense<8> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [134217728] ; Constancy: [8] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [134217728], constancy = [8], constant_value = <none> %4 = arith.divsi %0, %3 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [1], constant_value = <none> %5 = arith.andi %0, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [1], constant_value = <none> %6 = arith.ori %0, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [1], constant_value = <none> %7 = arith.xori %0, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [8] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [8], constant_value = <none> %8 = arith.andi %2, %4 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [8] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [8], constant_value = <none> %9 = arith.ori %2, %4 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [8] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [8], constant_value = <none> %10 = arith.xori %2, %4 : tensor<128xi32> return } // ----- -// CHECK-LABEL: select-func @select() {- // CHECK: Contiguity: [128] ; Divisibility: [1073741824] ; Constancy: [1] ; ConstantValue: [None]+// CHECK-LABEL: @select+func.func @select() {+ // CHECK: contiguity = [128], divisibility = [1073741824], constancy = [1], constant_value = <none> %0 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [4611686018427387904] ; Constancy: [128] ; ConstantValue: [0]+ // CHECK-NEXT: contiguity = [1], divisibility = [4611686018427387904], constancy = [128], constant_value = 0 %1 = arith.constant dense<0> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [128] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [128], constant_value = <none> %2 = arith.cmpi eq, %0, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [128] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [128], constant_value = <none> %3 = arith.cmpi slt, %0, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [4611686018427387904] ; Constancy: [1] ; ConstantValue: [0]+ // CHECK-NEXT: contiguity = [1], divisibility = [4611686018427387904], constancy = [1], constant_value = 0 %4 = arith.constant 0 : i1- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [4611686018427387904] ; Constancy: [128] ; ConstantValue: [0]+ // CHECK-NEXT: contiguity = [1], divisibility = [4611686018427387904], constancy = [128], constant_value = 0 %7 = tt.splat %4 : (i1) -> tensor<128xi1>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [4611686018427387904] ; Constancy: [128] ; ConstantValue: [0]- %5 = select %4, %3, %7 : tensor<128xi1>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [128] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [4611686018427387904], constancy = [128], constant_value = 0+ %5 = arith.select %4, %3, %7 : tensor<128xi1>+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [128], constant_value = <none> %8 = "triton_gpu.select"(%7, %3, %2) : (tensor<128xi1>, tensor<128xi1>, tensor<128xi1>) -> tensor<128xi1> return } // ----- -func @shift() {- // CHECK: Contiguity: [128] ; Divisibility: [1073741824] ; Constancy: [1] ; ConstantValue: [None]+func.func @shift() {+ // CHECK: contiguity = [128], divisibility = [1073741824], constancy = [1], constant_value = <none> %0 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [8] ; Constancy: [128] ; ConstantValue: [8]+ // CHECK-NEXT: contiguity = [1], divisibility = [8], constancy = [128], constant_value = 8 %1 = arith.constant dense<8> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [4] ; Constancy: [128] ; ConstantValue: [4]+ // CHECK-NEXT: contiguity = [1], divisibility = [4], constancy = [128], constant_value = 4 %2 = arith.constant dense<4> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [274877906944] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [274877906944], constancy = [1], constant_value = <none> %3 = arith.shli %0, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [67108864] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [67108864], constancy = [1], constant_value = <none> %4 = arith.shrsi %0, %2 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [128] ; Constancy: [128] ; ConstantValue: [128]+ // CHECK-NEXT: contiguity = [1], divisibility = [128], constancy = [128], constant_value = 128 %5 = arith.shli %1, %2 : tensor<128xi32> return } // ----- -func @max_min() {- // CHECK: Contiguity: [128] ; Divisibility: [1073741824] ; Constancy: [1] ; ConstantValue: [None]+func.func @max_min() {+ // CHECK: contiguity = [128], divisibility = [1073741824], constancy = [1], constant_value = <none> %0 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32>- // CHECK-NEXT: Contiguity: [128] ; Divisibility: [64] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [128], divisibility = [64], constancy = [1], constant_value = <none> %1 = tt.make_range {end = 192 : i32, start = 64 : i32} : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [1], constant_value = <none> %2 = arith.maxsi %0, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [1], constant_value = <none> %3 = arith.minsi %0, %1 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [8] ; Constancy: [128] ; ConstantValue: [8]+ // CHECK-NEXT: contiguity = [1], divisibility = [8], constancy = [128], constant_value = 8 %4 = arith.constant dense<8> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [4] ; Constancy: [128] ; ConstantValue: [4]+ // CHECK-NEXT: contiguity = [1], divisibility = [4], constancy = [128], constant_value = 4 %5 = arith.constant dense<4> : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [1] ; ConstantValue: [8]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [1], constant_value = 8 %6 = arith.maxsi %4, %5 : tensor<128xi32> return } // ----- -// CHECK-LABEL: for-func @for() {- // CHECK: Contiguity: [1, 1] ; Divisibility: [4611686018427387904, 4611686018427387904] ; Constancy: [128, 32] ; ConstantValue: [0]+// CHECK-LABEL: @for+func.func @for() {+ // CHECK: contiguity = [1, 1], divisibility = [4611686018427387904, 4611686018427387904], constancy = [128, 32], constant_value = 0 %a_init = arith.constant dense<0> : tensor<128x32xi32>- // CHECK-NEXT: Contiguity: [1, 1] ; Divisibility: [1, 1] ; Constancy: [128, 32] ; ConstantValue: [1]+ // CHECK-NEXT: contiguity = [1, 1], divisibility = [1, 1], constancy = [128, 32], constant_value = 1 %b_init = arith.constant dense<1> : tensor<128x32xi32>- // CHECK-NEXT: Contiguity: [1, 1] ; Divisibility: [4, 4] ; Constancy: [128, 32] ; ConstantValue: [4]+ // CHECK-NEXT: contiguity = [1, 1], divisibility = [4, 4], constancy = [128, 32], constant_value = 4 %c_init = arith.constant dense<4> : tensor<128x32xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [128] ; Constancy: [1] ; ConstantValue: [128]+ // CHECK-NEXT: contiguity = [1], divisibility = [128], constancy = [1], constant_value = 128 %ub = arith.constant 128 : index- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [4611686018427387904] ; Constancy: [1] ; ConstantValue: [0]+ // CHECK-NEXT: contiguity = [1], divisibility = [4611686018427387904], constancy = [1], constant_value = 0 %lb = arith.constant 0 : index- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [16] ; Constancy: [1] ; ConstantValue: [16]+ // CHECK-NEXT: contiguity = [1], divisibility = [16], constancy = [1], constant_value = 16 %step = arith.constant 16 : index %a, %b, %c = scf.for %iv = %lb to %ub step %step iter_args(%a = %a_init, %b = %b_init, %c = %c_init) -> (tensor<128x32xi32>, tensor<128x32xi32>, tensor<128x32xi32>) {- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [16] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [16], constancy = [1], constant_value = <none> %t = arith.index_cast %iv : index to i32- // CHECK: Contiguity: [1, 1] ; Divisibility: [1, 1] ; Constancy: [128, 32] ; ConstantValue: [None]- // CHECK: Contiguity: [1, 1] ; Divisibility: [1, 1] ; Constancy: [128, 32] ; ConstantValue: [None]- // CHECK: Contiguity: [1, 1] ; Divisibility: [4, 4] ; Constancy: [128, 32] ; ConstantValue: [4]+ // CHECK: contiguity = [1, 1], divisibility = [1, 1], constancy = [128, 32], constant_value = <none>+ // CHECK: contiguity = [1, 1], divisibility = [1, 1], constancy = [128, 32], constant_value = <none>+ // CHECK: contiguity = [1, 1], divisibility = [4, 4], constancy = [128, 32], constant_value = 4 scf.yield %b, %a, %c : tensor<128x32xi32>, tensor<128x32xi32>, tensor<128x32xi32> } return@@ -290,53 +290,53 @@ func @for() { // ----- -// CHECK-LABEL: permute_2d-func @permute_2d(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: i32 {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}) {- // CHECK: Contiguity: [1, 1] ; Divisibility: [1, 1] ; Constancy: [128, 128] ; ConstantValue: [1]+// CHECK-LABEL: @permute_2d+func.func @permute_2d(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: i32 {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}) {+ // CHECK: contiguity = [1, 1], divisibility = [1, 1], constancy = [128, 128], constant_value = 1 %cst = arith.constant dense<true> : tensor<128x128xi1>- // CHECK-NEXT: Contiguity: [1, 1] ; Divisibility: [1, 1] ; Constancy: [1, 1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1, 1], divisibility = [1, 1], constancy = [1, 1], constant_value = <none> %cst_0 = arith.constant dense<0.000000e+00> : tensor<128x128xf32>- // CHECK-NEXT: Contiguity: [128] ; Divisibility: [1073741824] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [128], divisibility = [1073741824], constancy = [1], constant_value = <none> %0 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32>- // CHECK-NEXT: Contiguity: [128] ; Divisibility: [1073741824] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [128], divisibility = [1073741824], constancy = [1], constant_value = <none> %1 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32>- // CHECK-NEXT: Contiguity: [128, 1] ; Divisibility: [1073741824, 1] ; Constancy: [1, 1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [128, 1], divisibility = [1073741824, 1], constancy = [1, 1], constant_value = <none> %2 = tt.expand_dims %0 {axis = 1 : i32} : (tensor<128xi32>) -> tensor<128x1xi32>- // CHECK-NEXT: Contiguity: [1, 1] ; Divisibility: [16, 16] ; Constancy: [128, 1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1, 1], divisibility = [16, 16], constancy = [128, 1], constant_value = <none> %3 = tt.splat %arg1 : (i32) -> tensor<128x1xi32>- // CHECK-NEXT: Contiguity: [1, 1] ; Divisibility: [17179869184, 16] ; Constancy: [1, 1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1, 1], divisibility = [17179869184, 16], constancy = [1, 1], constant_value = <none> %4 = arith.muli %2, %3 : tensor<128x1xi32>- // CHECK-NEXT: Contiguity: [1, 1] ; Divisibility: [16, 16] ; Constancy: [128, 1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1, 1], divisibility = [16, 16], constancy = [128, 1], constant_value = <none> %5 = tt.splat %arg0 : (!tt.ptr<f32>) -> tensor<128x1x!tt.ptr<f32>>- // CHECK-NEXT: Contiguity: [1, 1] ; Divisibility: [16, 16] ; Constancy: [1, 1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1, 1], divisibility = [16, 16], constancy = [1, 1], constant_value = <none> %6 = tt.addptr %5, %4 : tensor<128x1x!tt.ptr<f32>>, tensor<128x1xi32>- // CHECK-NEXT: Contiguity: [1, 128] ; Divisibility: [1, 1073741824] ; Constancy: [1, 1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1, 128], divisibility = [1, 1073741824], constancy = [1, 1], constant_value = <none> %7 = tt.expand_dims %1 {axis = 0 : i32}: (tensor<128xi32>) -> tensor<1x128xi32>- // CHECK-NEXT: Contiguity: [1, 1] ; Divisibility: [16, 16] ; Constancy: [1, 128] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1, 1], divisibility = [16, 16], constancy = [1, 128], constant_value = <none> %8 = tt.broadcast %6 : (tensor<128x1x!tt.ptr<f32>>) -> tensor<128x128x!tt.ptr<f32>>- // CHECK-NEXT: Contiguity: [1, 128] ; Divisibility: [1, 1073741824] ; Constancy: [128, 1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1, 128], divisibility = [1, 1073741824], constancy = [128, 1], constant_value = <none> %9 = tt.broadcast %7 : (tensor<1x128xi32>) -> tensor<128x128xi32>- // CHECK-NEXT: Contiguity: [1, 128] ; Divisibility: [1, 16] ; Constancy: [1, 1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1, 128], divisibility = [1, 16], constancy = [1, 1], constant_value = <none> %10 = tt.addptr %8, %9 : tensor<128x128x!tt.ptr<f32>>, tensor<128x128xi32>- // CHECK-NEXT: Contiguity: [128, 1] ; Divisibility: [1073741824, 1] ; Constancy: [1, 1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [128, 1], divisibility = [1073741824, 1], constancy = [1, 1], constant_value = <none> %11 = tt.expand_dims %0 {axis = 1 : i32}: (tensor<128xi32>) -> tensor<128x1xi32>- // CHECK-NEXT: Contiguity: [1, 1] ; Divisibility: [16, 16] ; Constancy: [128, 1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1, 1], divisibility = [16, 16], constancy = [128, 1], constant_value = <none> %12 = tt.splat %arg2 : (!tt.ptr<f32>) -> tensor<128x1x!tt.ptr<f32>>- // CHECK-NEXT: Contiguity: [128, 1] ; Divisibility: [16, 1] ; Constancy: [1, 1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [128, 1], divisibility = [16, 1], constancy = [1, 1], constant_value = <none> %13 = tt.addptr %12, %11 : tensor<128x1x!tt.ptr<f32>>, tensor<128x1xi32>- // CHECK-NEXT: Contiguity: [1, 128] ; Divisibility: [1, 1073741824] ; Constancy: [1, 1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1, 128], divisibility = [1, 1073741824], constancy = [1, 1], constant_value = <none> %14 = tt.expand_dims %1 {axis = 0 : i32} : (tensor<128xi32>) -> tensor<1x128xi32>- // CHECK-NEXT: Contiguity: [1, 1] ; Divisibility: [16, 16] ; Constancy: [1, 128] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1, 1], divisibility = [16, 16], constancy = [1, 128], constant_value = <none> %15 = tt.splat %arg3 : (i32) -> tensor<1x128xi32>- // CHECK-NEXT: Contiguity: [1, 1] ; Divisibility: [16, 17179869184] ; Constancy: [1, 1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1, 1], divisibility = [16, 17179869184], constancy = [1, 1], constant_value = <none> %16 = arith.muli %14, %15 : tensor<1x128xi32>- // CHECK-NEXT: Contiguity: [128, 1] ; Divisibility: [16, 1] ; Constancy: [1, 128] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [128, 1], divisibility = [16, 1], constancy = [1, 128], constant_value = <none> %17 = tt.broadcast %13 : (tensor<128x1x!tt.ptr<f32>>) -> tensor<128x128x!tt.ptr<f32>>- // CHECK-NEXT: Contiguity: [1, 1] ; Divisibility: [16, 17179869184] ; Constancy: [128, 1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1, 1], divisibility = [16, 17179869184], constancy = [128, 1], constant_value = <none> %18 = tt.broadcast %16 : (tensor<1x128xi32>) -> tensor<128x128xi32>- // CHECK-NEXT: Contiguity: [128, 1] ; Divisibility: [16, 1] ; Constancy: [1, 1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [128, 1], divisibility = [16, 1], constancy = [1, 1], constant_value = <none> %19 = tt.addptr %17, %18 : tensor<128x128x!tt.ptr<f32>>, tensor<128x128xi32>- // CHECK-NEXT: Contiguity: [1, 1] ; Divisibility: [1, 1] ; Constancy: [1, 1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1, 1], divisibility = [1, 1], constancy = [1, 1], constant_value = <none> %20 = tt.load %10, %cst, %cst_0 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<128x128xf32> tt.store %19, %20, %cst : tensor<128x128xf32> return@@ -347,29 +347,29 @@ func @permute_2d(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: i32 {t module { // This is a tiny test for verifying StoreOp-related alignment, It simply store a constant to a buffer.-// CHECK-LABEL: store_constant_align-func @store_constant_align(%addr: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %n: i32 {tt.divisibility = 16 : i32}) {- // CHECK: Contiguity: [1] ; Divisibility: [1] ; Constancy: [1] ; ConstantValue: [None]+// CHECK-LABEL: @store_constant_align+func.func @store_constant_align(%addr: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %n: i32 {tt.divisibility = 16 : i32}) {+ // CHECK: contiguity = [1], divisibility = [1], constancy = [1], constant_value = <none> %pid = tt.get_program_id {axis = 0 : i32} : i32- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [128] ; Constancy: [1] ; ConstantValue: [128]+ // CHECK-NEXT: contiguity = [1], divisibility = [128], constancy = [1], constant_value = 128 %c128_i32 = arith.constant 128 : i32- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [128] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [128], constancy = [1], constant_value = <none> %1 = arith.muli %pid, %c128_i32 : i32- // CHECK-NEXT: Contiguity: [128] ; Divisibility: [1073741824] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [128], divisibility = [1073741824], constancy = [1], constant_value = <none> %2 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [128] ; Constancy: [128] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [128], constancy = [128], constant_value = <none> %3 = tt.splat %1 : (i32) -> tensor<128xi32>- // CHECK-NEXT: Contiguity: [128] ; Divisibility: [128] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [128], divisibility = [128], constancy = [1], constant_value = <none> %4 = arith.addi %3, %2 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [16] ; Constancy: [128] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [16], constancy = [128], constant_value = <none> %5 = tt.splat %addr : (!tt.ptr<f32>) -> tensor<128x!tt.ptr<f32>>- // CHECK-NEXT: Contiguity: [128] ; Divisibility: [16] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [128], divisibility = [16], constancy = [1], constant_value = <none> %6 = tt.addptr %5, %4 : tensor<128x!tt.ptr<f32>>, tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [16] ; Constancy: [128] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [16], constancy = [128], constant_value = <none> %9 = tt.splat %n : (i32) -> tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [16] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [16], constant_value = <none> %mask = arith.cmpi slt, %4, %9 : tensor<128xi32>- // CHECK-NEXT: Contiguity: [1] ; Divisibility: [1] ; Constancy: [1] ; ConstantValue: [None]+ // CHECK-NEXT: contiguity = [1], divisibility = [1], constancy = [1], constant_value = <none> %cst = arith.constant dense<0.0> : tensor<128xf32> tt.store %5, %cst, %mask : tensor<128xf32> return@@ -381,8 +381,8 @@ func @store_constant_align(%addr: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %n: // This IR is dumped from vecadd test. // Note, the hint {tt.divisibility = 16 : i32} for %n_elements affects the alignment of mask.-// CHECK-LABEL: vecadd_mask_align_16-func @vecadd_mask_align_16(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %n_elements: i32 {tt.divisibility = 16 : i32}) {+// CHECK-LABEL: @vecadd_mask_align_16+func.func @vecadd_mask_align_16(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %n_elements: i32 {tt.divisibility = 16 : i32}) { %c64_i32 = arith.constant 64 : i32 %0 = tt.get_program_id {axis = 0 : i32} : i32 %1 = arith.muli %0, %c64_i32 : i32@@ -394,13 +394,13 @@ func @vecadd_mask_align_16(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %ar %7 = tt.splat %arg1 : (!tt.ptr<f32>) -> tensor<64x!tt.ptr<f32>> %8 = tt.addptr %7, %4 : tensor<64x!tt.ptr<f32>>, tensor<64xi32> %9 = tt.splat %n_elements : (i32) -> tensor<64xi32>- // CHECK: Contiguity: [1] ; Divisibility: [1] ; Constancy: [16] ; ConstantValue: [None] ( %{{.*}} = arith.cmpi slt, %{{.*}}, %{{.*}} : tensor<64xi32> )+ // CHECK: arith.cmpi slt, %{{.*}} => contiguity = [1], divisibility = [1], constancy = [16], constant_value = <none> %mask = arith.cmpi slt, %4, %9 : tensor<64xi32> %11 = tt.load %6, %mask {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<64xf32> %12 = tt.load %8, %mask {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<64xf32> %13 = arith.addf %11, %12 : tensor<64xf32> %14 = tt.splat %arg2 : (!tt.ptr<f32>) -> tensor<64x!tt.ptr<f32>>- // CHECK: Contiguity: [64] ; Divisibility: [16] ; Constancy: [1] ; ConstantValue: [None] ( %{{.*}} = tt.addptr %{{.*}}, %{{.*}} : tensor<64x!tt.ptr<f32>>, tensor<64xi32> )+ // CHECK: tt.addptr %{{.*}} => contiguity = [64], divisibility = [16], constancy = [1], constant_value = <none> %15 = tt.addptr %14, %4 : tensor<64x!tt.ptr<f32>>, tensor<64xi32> tt.store %15, %13, %mask : tensor<64xf32> return@@ -410,8 +410,8 @@ func @vecadd_mask_align_16(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %ar // This IR is dumped from vecadd test. // Note, there is no divisibility hint for %n_elements, Triton should assume its divisibility to be 1 by default.-// CHECK-LABEL: vecadd_mask_align_1-func @vecadd_mask_align_1(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %n_elements: i32) {+// CHECK-LABEL: @vecadd_mask_align_1+func.func @vecadd_mask_align_1(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %n_elements: i32) { %c64_i32 = arith.constant 64 : i32 %0 = tt.get_program_id {axis = 0 : i32} : i32 %1 = arith.muli %0, %c64_i32 : i32@@ -423,7 +423,7 @@ func @vecadd_mask_align_1(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg %7 = tt.splat %arg1 : (!tt.ptr<f32>) -> tensor<64x!tt.ptr<f32>> %8 = tt.addptr %7, %4 : tensor<64x!tt.ptr<f32>>, tensor<64xi32> %9 = tt.splat %n_elements : (i32) -> tensor<64xi32>- // CHECK: Contiguity: [1] ; Divisibility: [1] ; Constancy: [1] ; ConstantValue: [None] ( %{{.*}} = arith.cmpi slt, %{{.*}}, %{{.*}} : tensor<64xi32> )+ // CHECK: arith.cmpi slt, %{{.*}} => contiguity = [1], divisibility = [1], constancy = [1], constant_value = <none> %10 = arith.cmpi slt, %4, %9 : tensor<64xi32> %11 = tt.load %6, %10 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<64xf32> %12 = tt.load %8, %10 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<64xf32>diff --git a/test/Analysis/test-allocation.mlir b/test/Analysis/test-allocation.mlirindex efb00c404d..f79222aa7b 100644--- a/test/Analysis/test-allocation.mlir+++ b/test/Analysis/test-allocation.mlir@@ -13,7 +13,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: matmul_loop-func @matmul_loop(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) {+func.func @matmul_loop(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) { %a_ptr_init = tt.broadcast %A : (!tt.ptr<f16>) -> tensor<128x32x!tt.ptr<f16>, #AL> %b_ptr_init = tt.broadcast %B : (!tt.ptr<f16>) -> tensor<32x128x!tt.ptr<f16>, #BL> @@ -46,7 +46,7 @@ func @matmul_loop(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B // Shared memory is available after a tensor's liveness range ends // CHECK-LABEL: reusable-func @reusable(%A : !tt.ptr<f16>) {+func.func @reusable(%A : !tt.ptr<f16>) { %cst1 = arith.constant dense<true> : tensor<128x32xi1, #AL> %cst2 = arith.constant dense<0.000000e+00> : tensor<128x32xf16, #AL> %cst3 = arith.constant dense<true> : tensor<32x128xi1, #AL>@@ -78,7 +78,7 @@ func @reusable(%A : !tt.ptr<f16>) { // %cst1->%cst4 // %cst3->%g->%h->%i // CHECK-LABEL: preallocate-func @preallocate(%A : !tt.ptr<f16>) {+func.func @preallocate(%A : !tt.ptr<f16>) { // CHECK: offset = 0, size = 512 %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> // CHECK-NEXT: offset = 1024, size = 512@@ -113,7 +113,7 @@ func @preallocate(%A : !tt.ptr<f16>) { // Unused tensors are immediately released // CHECK-LABEL: unused-func @unused(%A : !tt.ptr<f16>) {+func.func @unused(%A : !tt.ptr<f16>) { // CHECK: offset = 0, size = 1024 %cst0 = arith.constant dense<0.000000e+00> : tensor<32x16xf16, #A_SHARED> // CHECK-NEXT: offset = 0, size = 512@@ -128,7 +128,7 @@ func @unused(%A : !tt.ptr<f16>) { // cst0 is alive through the entire function, it cannot be released before the end of the function // CHECK-LABEL: longlive-func @longlive(%A : !tt.ptr<f16>) {+func.func @longlive(%A : !tt.ptr<f16>) { // CHECK: offset = 0, size = 512 %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> // CHECK-NEXT: offset = 512, size = 512@@ -156,7 +156,7 @@ func @longlive(%A : !tt.ptr<f16>) { } // CHECK-LABEL: alloc-func @alloc(%A : !tt.ptr<f16>) {+func.func @alloc(%A : !tt.ptr<f16>) { // CHECK: offset = 0, size = 512 %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> %cst1 = arith.constant dense<0.000000e+00> : tensor<16x32xf16, #AL>@@ -167,7 +167,7 @@ func @alloc(%A : !tt.ptr<f16>) { } // CHECK-LABEL: scratch-func @scratch() {+func.func @scratch() { %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #AL> // CHECK: scratch offset = 0, size = 512 %b = tt.reduce %cst0 {redOp = 1 : i32, axis = 0 : i32} : tensor<16x16xf16, #AL> -> tensor<16xf16, #sliceAd0>@@ -176,7 +176,7 @@ func @scratch() { } // CHECK-LABEL: trans-func @trans(%A : !tt.ptr<f16>) {+func.func @trans(%A : !tt.ptr<f16>) { // CHECK: offset = 0, size = 1024 %tensor = arith.constant dense<0.000000e+00> : tensor<16x32xf16, #A_SHARED> %b = tt.trans %tensor : (tensor<16x32xf16, #A_SHARED>) -> tensor<32x16xf16, #A_SHARED_T>@@ -184,7 +184,7 @@ func @trans(%A : !tt.ptr<f16>) { } // CHECK-LABEL: insert_slice_async-func @insert_slice_async(%A : !tt.ptr<f16>, %i1 : i1) {+func.func @insert_slice_async(%A : !tt.ptr<f16>, %i1 : i1) { %a_ptr = tt.broadcast %A : (!tt.ptr<f16>) -> tensor<16x16x!tt.ptr<f16>, #AL> %mask = tt.splat %i1 : (i1) -> tensor<16x16xi1, #AL> %other = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #AL>@@ -197,7 +197,7 @@ func @insert_slice_async(%A : !tt.ptr<f16>, %i1 : i1) { } // CHECK-LABEL: extract_slice-func @extract_slice(%A : !tt.ptr<f16>) {+func.func @extract_slice(%A : !tt.ptr<f16>) { // CHECK: offset = 0, size = 512 %cst0 = arith.constant dense<0.000000e+00> : tensor<1x16x16xf16, #A_SHARED> %index = arith.constant 0 : index@@ -209,7 +209,7 @@ func @extract_slice(%A : !tt.ptr<f16>) { // B0 -> (B1) -> B0 // Memory used by B1 can be reused by B0. // CHECK-LABEL: if-func @if(%i1 : i1) {+func.func @if(%i1 : i1) { // CHECK: offset = 0, size = 512 %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> // CHECK-NEXT: offset = 512, size = 512@@ -233,7 +233,7 @@ func @if(%i1 : i1) { // B0 -> (B1) -> (B2) -> B0 // Memory used by B0 cannot be reused by B1 or B2. // CHECK-LABEL: if_else-func @if_else(%i1 : i1) {+func.func @if_else(%i1 : i1) { // CHECK: offset = 0, size = 512 %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> // CHECK-NEXT: offset = 512, size = 512@@ -260,7 +260,7 @@ func @if_else(%i1 : i1) { // Block arguments and yields are memory aliases that do not trigger a new // allocation. // CHECK-LABEL: for-func @for(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) {+func.func @for(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) { // CHECK: offset = 0, size = 8192 %a_shared_init = arith.constant dense<0.00e+00> : tensor<128x32xf16, #A_SHARED> // CHECK-NEXT: offset = 8192, size = 8192@@ -275,7 +275,7 @@ func @for(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.p } // CHECK-LABEL: for_if_slice-func @for_if_slice(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>, %i1 : i1) {+func.func @for_if_slice(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>, %i1 : i1) { // CHECK: offset = 0, size = 8192 %a_shared_init = arith.constant dense<0.00e+00> : tensor<128x32xf16, #A_SHARED> // CHECK-NEXT: offset = 8192, size = 8192@@ -296,7 +296,7 @@ func @for_if_slice(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, % // c0 cannot be released in the loop // CHECK-LABEL: for_use_ancestor-func @for_use_ancestor(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>, %i1 : i1) {+func.func @for_use_ancestor(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>, %i1 : i1) { // CHECK: offset = 0, size = 8192 %a_shared_init = arith.constant dense<0.00e+00> : tensor<128x32xf16, #A_SHARED> // CHECK-NEXT: offset = 8192, size = 8192@@ -316,7 +316,7 @@ func @for_use_ancestor(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16 // a_shared_init, b_shared_init, and c_shared_init's liveness ranges are span over the entire function before cst2. // So they cannot be reused by cst0 and cst1, but can be reused by cst2. // CHECK-LABEL: for_if_for-func @for_if_for(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>, %i1 : i1) {+func.func @for_if_for(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>, %i1 : i1) { // CHECK: offset = 0, size = 8192 %a_shared_init = arith.constant dense<0.00e+00> : tensor<128x32xf16, #A_SHARED> // CHECK-NEXT: offset = 8192, size = 8192diff --git a/test/Analysis/test-membar.mlir b/test/Analysis/test-membar.mlirindex 7199e5f53d..17880b2094 100644--- a/test/Analysis/test-membar.mlir+++ b/test/Analysis/test-membar.mlir@@ -14,7 +14,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: matmul_loop // There shouldn't be any membar with the dot op encoding.-func @matmul_loop(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) {+func.func @matmul_loop(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) { %a_ptr_init = tt.broadcast %A : (!tt.ptr<f16>) -> tensor<128x32x!tt.ptr<f16>, #AL> %b_ptr_init = tt.broadcast %B : (!tt.ptr<f16>) -> tensor<32x128x!tt.ptr<f16>, #BL> @@ -42,7 +42,7 @@ func @matmul_loop(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B } // CHECK-LABEL: raw_single_block-func @raw_single_block(%A : !tt.ptr<f16>) {+func.func @raw_single_block(%A : !tt.ptr<f16>) { %cst1 = arith.constant dense<true> : tensor<128x32xi1, #AL> %cst2 = arith.constant dense<0.000000e+00> : tensor<128x32xf16, #AL> %a_ptr = tt.broadcast %A : (!tt.ptr<f16>) -> tensor<128x32x!tt.ptr<f16>, #AL>@@ -54,7 +54,7 @@ func @raw_single_block(%A : !tt.ptr<f16>) { } // CHECK-LABEL: war_single_block-func @war_single_block(%A : !tt.ptr<f16>) {+func.func @war_single_block(%A : !tt.ptr<f16>) { %cst1 = arith.constant dense<true> : tensor<128x32xi1, #AL> %cst2 = arith.constant dense<0.000000e+00> : tensor<128x32xf16, #AL> %a_ptr = tt.broadcast %A : (!tt.ptr<f16>) -> tensor<128x32x!tt.ptr<f16>, #AL>@@ -70,7 +70,7 @@ func @war_single_block(%A : !tt.ptr<f16>) { } // CHECK-LABEL: scratch-func @scratch() {+func.func @scratch() { %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> // CHECK: Membar 1 %a = tt.cat %cst0, %cst0 {axis = 0} : (tensor<16x16xf16, #A_SHARED>, tensor<16x16xf16, #A_SHARED>) -> tensor<32x16xf16, #A_SHARED>@@ -81,7 +81,7 @@ func @scratch() { } // CHECK-LABEL: async_wait-func @async_wait() {+func.func @async_wait() { %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> // CHECK: Membar 1 %a = tt.cat %cst0, %cst0 {axis = 0} : (tensor<16x16xf16, #A_SHARED>, tensor<16x16xf16, #A_SHARED>) -> tensor<32x16xf16, #A_SHARED>@@ -92,7 +92,7 @@ func @async_wait() { } // CHECK-LABEL: alloc-func @alloc() {+func.func @alloc() { %cst0 = triton_gpu.alloc_tensor : tensor<16x16xf16, #A_SHARED> %a = tt.cat %cst0, %cst0 {axis = 0} : (tensor<16x16xf16, #A_SHARED>, tensor<16x16xf16, #A_SHARED>) -> tensor<32x16xf16, #A_SHARED> // CHECK: Membar 2@@ -101,7 +101,7 @@ func @alloc() { } // CHECK-LABEL: extract_slice-func @extract_slice() {+func.func @extract_slice() { %cst0 = arith.constant dense<0.000000e+00> : tensor<1x16x16xf16, #A_SHARED> %index = arith.constant 0 : index %cst1 = tensor.extract_slice %cst0[%index, 0, 0][1, 16, 16][1, 1, 1] : tensor<1x16x16xf16, #A_SHARED> to tensor<16x16xf16, #A_SHARED>@@ -113,14 +113,14 @@ func @extract_slice() { } // CHECK-LABEL: trans-func @trans() {+func.func @trans() { %cst0 = arith.constant dense<0.000000e+00> : tensor<16x32xf16, #A_SHARED> %b = tt.trans %cst0 : (tensor<16x32xf16, #A_SHARED>) -> tensor<32x16xf16, #A_SHARED_T> return } // CHECK-LABEL: insert_slice_async-func @insert_slice_async(%A : !tt.ptr<f16>, %i1 : i1) {+func.func @insert_slice_async(%A : !tt.ptr<f16>, %i1 : i1) { %a_ptr = tt.broadcast %A : (!tt.ptr<f16>) -> tensor<16x16x!tt.ptr<f16>, #AL> %mask = tt.splat %i1 : (i1) -> tensor<16x16xi1, #AL> %other = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #AL>@@ -135,7 +135,7 @@ func @insert_slice_async(%A : !tt.ptr<f16>, %i1 : i1) { } // CHECK-LABEL: insert_slice-func @insert_slice(%A : !tt.ptr<f16>, %i1 : i1) {+func.func @insert_slice(%A : !tt.ptr<f16>, %i1 : i1) { %a_ptr = tt.broadcast %A : (!tt.ptr<f16>) -> tensor<16x16x!tt.ptr<f16>, #AL> %mask = tt.splat %i1 : (i1) -> tensor<16x16xi1, #AL> %other = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #AL>@@ -153,7 +153,7 @@ func @insert_slice(%A : !tt.ptr<f16>, %i1 : i1) { // If branch inserted a barrier for %cst0 and %cst1, but else didn't, then the barrier should be inserted in the parent region // CHECK-LABEL: multi_blocks-func @multi_blocks(%i1 : i1) {+func.func @multi_blocks(%i1 : i1) { %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> %cst1 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> scf.if %i1 {@@ -174,7 +174,7 @@ func @multi_blocks(%i1 : i1) { // Both branches inserted a barrier for %cst0 and %cst1, then the barrier doesn't need to be inserted in the parent region // CHECK-LABEL: multi_blocks_join_barrier-func @multi_blocks_join_barrier(%i1 : i1) {+func.func @multi_blocks_join_barrier(%i1 : i1) { %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> %cst1 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> scf.if %i1 {@@ -192,7 +192,7 @@ func @multi_blocks_join_barrier(%i1 : i1) { // Read yielded tensor requires a barrier // CHECK-LABEL: multi_blocks_yield-func @multi_blocks_yield(%i1 : i1) {+func.func @multi_blocks_yield(%i1 : i1) { %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> %cst1 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> %a = scf.if %i1 -> (tensor<32x16xf16, #A_SHARED>) {@@ -212,7 +212,7 @@ func @multi_blocks_yield(%i1 : i1) { // Conservatively add a barrier as if the branch (%i1) is never taken // CHECK-LABEL: multi_blocks_noelse-func @multi_blocks_noelse(%i1 : i1) {+func.func @multi_blocks_noelse(%i1 : i1) { %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> %cst1 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> scf.if %i1 {@@ -226,7 +226,7 @@ func @multi_blocks_noelse(%i1 : i1) { // Conservatively add a barrier as if the branch (%i2) is never taken // CHECK-LABEL: multi_blocks_nested_scf-func @multi_blocks_nested_scf(%i1 : i1, %i2 : i1) {+func.func @multi_blocks_nested_scf(%i1 : i1, %i2 : i1) { %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> %cst1 = arith.constant dense<0.000000e+00> : tensor<16x16xf16, #A_SHARED> scf.if %i1 {@@ -247,7 +247,7 @@ func @multi_blocks_nested_scf(%i1 : i1, %i2 : i1) { } // CHECK-LABEL: for-func @for(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) {+func.func @for(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) { %a_shared_init = arith.constant dense<0.00e+00> : tensor<128x32xf16, #A_SHARED> %b_shared_init = arith.constant dense<0.00e+00> : tensor<128x32xf16, #A_SHARED> %c_shared_init = arith.constant dense<0.00e+00> : tensor<128x32xf16, #A_SHARED>@@ -262,7 +262,7 @@ func @for(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.p // Although a_shared and b_shared are synced before entering the loop, // they are reassociated with aliases (c_shared) and thus require a barrier. // CHECK-LABEL: for_alias-func @for_alias(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) {+func.func @for_alias(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) { %a_shared_init = arith.constant dense<0.00e+00> : tensor<128x32xf16, #A_SHARED> %b_shared_init = arith.constant dense<0.00e+00> : tensor<128x32xf16, #A_SHARED> // CHECK-NEXT: Membar 2@@ -282,7 +282,7 @@ func @for_alias(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : // Although cst2 is not an argument of scf.yield, its memory is reused by cst1. // So we need a barrier both before and after cst1 // CHECK-LABEL: for_reuse-func @for_reuse(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) {+func.func @for_reuse(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) { %a_shared_init = arith.constant dense<0.00e+00> : tensor<128x32xf16, #A_SHARED> %b_shared_init = arith.constant dense<0.00e+00> : tensor<128x32xf16, #A_SHARED> // CHECK-NEXT: Membar 2@@ -302,7 +302,7 @@ func @for_reuse(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : // CHECK-LABEL: for_reuse_nested-func @for_reuse_nested(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) {+func.func @for_reuse_nested(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) { %a_shared_init = arith.constant dense<0.00e+00> : tensor<128x32xf16, #A_SHARED> %b_shared_init = arith.constant dense<0.00e+00> : tensor<128x32xf16, #A_SHARED> // CHECK-NEXT: Membar 2diff --git a/test/Conversion/triton_ops.mlir b/test/Conversion/triton_ops.mlirindex e9ee502435..0e979b148d 100644--- a/test/Conversion/triton_ops.mlir+++ b/test/Conversion/triton_ops.mlir@@ -1,6 +1,6 @@ // RUN: triton-opt %s | FileCheck %s -func @cast_ops(%scalar_ptr: !tt.ptr<f32>, %scalar_f32: f32, %scalar_i64: i64) {+func.func @cast_ops(%scalar_ptr: !tt.ptr<f32>, %scalar_f32: f32, %scalar_i64: i64) { // scalar -> scalar // CHECK: i64 -> !tt.ptr<f32> %0 = tt.int_to_ptr %scalar_i64 : i64 -> !tt.ptr<f32>@@ -35,7 +35,7 @@ func @cast_ops(%scalar_ptr: !tt.ptr<f32>, %scalar_f32: f32, %scalar_i64: i64) { return } -func @addptr_ops(%scalar_ptr: !tt.ptr<f32>, %scalar_i32: i32) {+func.func @addptr_ops(%scalar_ptr: !tt.ptr<f32>, %scalar_i32: i32) { // scalar -> scalar // CHECK: !tt.ptr<f32> %0 = tt.addptr %scalar_ptr, %scalar_i32 : !tt.ptr<f32>, i32@@ -54,7 +54,7 @@ func @addptr_ops(%scalar_ptr: !tt.ptr<f32>, %scalar_i32: i32) { return } -func @load_store_ops_scalar(%ptr: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %mask : i1) {+func.func @load_store_ops_scalar(%ptr: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %mask : i1) { // Test if Load/Store ops can handle scalar values %other = arith.constant 0.0e+0 : f32 @@ -76,7 +76,7 @@ func @load_store_ops_scalar(%ptr: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %ma return } -func @reduce_ops_infer(%ptr: !tt.ptr<f32>, %v : tensor<1x2x4xf32>) {+func.func @reduce_ops_infer(%ptr: !tt.ptr<f32>, %v : tensor<1x2x4xf32>) { // Test if reduce ops infer types correctly // CHECK: %{{.*}} = tt.reduce %{{.*}} -> tensor<2x4xf32>@@ -101,7 +101,7 @@ func @reduce_ops_infer(%ptr: !tt.ptr<f32>, %v : tensor<1x2x4xf32>) { return } -func @dot_ops_infer(%ptr: !tt.ptr<f32>, %v : f32) {+func.func @dot_ops_infer(%ptr: !tt.ptr<f32>, %v : f32) { // Test if reduce ops infer types correctly %v128x32 = tt.splat %v : (f32) -> tensor<128x32xf32> %v32x128 = tt.splat %v : (f32) -> tensor<32x128xf32>diff --git a/test/Conversion/triton_to_tritongpu.mlir b/test/Conversion/triton_to_tritongpu.mlirindex a160bc8815..b461ca542f 100644--- a/test/Conversion/triton_to_tritongpu.mlir+++ b/test/Conversion/triton_to_tritongpu.mlir@@ -1,6 +1,6 @@ // RUN: triton-opt %s -split-input-file -convert-triton-to-tritongpu=num-warps=2 | FileCheck %s -func @ops() {+func.func @ops() { // CHECK: module attributes {"triton_gpu.num-warps" = 2 : i32} {{.*}} %a = arith.constant dense<1.00e+00> : tensor<128x32xf16> %b = arith.constant dense<2.00e+00> : tensor<32x128xf16>@@ -11,7 +11,7 @@ func @ops() { // ----- -func @load_ops(%ptr: !tt.ptr<f32> {tt.divisibility = 16 : i32}) {+func.func @load_ops(%ptr: !tt.ptr<f32> {tt.divisibility = 16 : i32}) { // Test if LoadOp is lowered properly (see #771) %ptrs = tt.splat %ptr : (!tt.ptr<f32>) -> tensor<128x!tt.ptr<f32>> %mask = arith.constant dense<true> : tensor<128xi1>@@ -30,7 +30,7 @@ func @load_ops(%ptr: !tt.ptr<f32> {tt.divisibility = 16 : i32}) { // ----- -func @reduce_ops(%ptr: !tt.ptr<f32> {tt.divisibility = 16 : i32}) {+func.func @reduce_ops(%ptr: !tt.ptr<f32> {tt.divisibility = 16 : i32}) { // Test if the total number of threadsPerWarp is 32 // Test if the total number of warps is 2 // CHECK: #blocked0 = #triton_gpu.blocked<{sizePerThread = [1, 1], threadsPerWarp = [4, 8], warpsPerCTA = [1, 2], order = [0, 1]}>diff --git a/test/Conversion/tritongpu_to_llvm.mlir b/test/Conversion/tritongpu_to_llvm.mlirindex e9e7d5a340..507b362c99 100644--- a/test/Conversion/tritongpu_to_llvm.mlir+++ b/test/Conversion/tritongpu_to_llvm.mlir@@ -4,7 +4,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK: llvm.func @test_empty_kernel(%arg0: i32, %arg1: !llvm.ptr<f16, 1>) // Here the 128 comes from the 4 in module attribute multiples 32 // CHECK: attributes {nvvm.kernel = 1 : ui1, nvvm.maxntid = 128 : i32} {{.*}}- func @test_empty_kernel(%lb : index, %A : !tt.ptr<f16>) {+ func.func @test_empty_kernel(%lb : index, %A : !tt.ptr<f16>) { // CHECK: llvm.return return }@@ -15,7 +15,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #blocked0 = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}> module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: basic_load- func @basic_load(%a_ptr_init : tensor<256x!tt.ptr<f32>, #blocked0>, %cst : tensor<256xi1, #blocked0>, %cst_0 : tensor<256xf32, #blocked0>) {+ func.func @basic_load(%a_ptr_init : tensor<256x!tt.ptr<f32>, #blocked0>, %cst : tensor<256xi1, #blocked0>, %cst_0 : tensor<256xf32, #blocked0>) { // CHECK: llvm.inline_asm // CHECK: llvm.inline_asm %1 = tt.load %a_ptr_init, %cst, %cst_0 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<256xf32, #blocked0>@@ -28,7 +28,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #blocked0 = #triton_gpu.blocked<{sizePerThread = [2], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}> module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: vectorized_load- func @vectorized_load(%a_ptr_init : tensor<256x!tt.ptr<f32>, #blocked0>, %cst : tensor<256xi1, #blocked0>, %cst_0 : tensor<256xf32, #blocked0>) {+ func.func @vectorized_load(%a_ptr_init : tensor<256x!tt.ptr<f32>, #blocked0>, %cst : tensor<256xi1, #blocked0>, %cst_0 : tensor<256xf32, #blocked0>) { // CHECK: llvm.inline_asm // CHECK-SAME: ld.global.b32 // CHECK: llvm.inline_asm@@ -43,7 +43,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #blocked0 = #triton_gpu.blocked<{sizePerThread = [8], threadsPerWarp = [32], warpsPerCTA = [1], order = [0]}> module attributes {"triton_gpu.num-warps" = 1 : i32} { // CHECK-LABEL: vectorized_load_f16- func @vectorized_load_f16(%a_ptr_init: tensor<256x!tt.ptr<f16>, #blocked0>, %cst : tensor<256xi1, #blocked0>, %cst_0 : tensor<256xf16, #blocked0>) {+ func.func @vectorized_load_f16(%a_ptr_init: tensor<256x!tt.ptr<f16>, #blocked0>, %cst : tensor<256xi1, #blocked0>, %cst_0 : tensor<256xf16, #blocked0>) { // CHECK: llvm.inline_asm // CHECK-SAME: ld.global.b16 // CHECK: llvm.inline_asm@@ -59,7 +59,7 @@ module attributes {"triton_gpu.num-warps" = 1 : i32} { #blocked0 = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [8], order = [0]}> module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: masked_load_const_other- func @masked_load_const_other(%a_ptr_init : tensor<256x!tt.ptr<f32>, #blocked0>, %cst : tensor<256xi1, #blocked0>) {+ func.func @masked_load_const_other(%a_ptr_init : tensor<256x!tt.ptr<f32>, #blocked0>, %cst : tensor<256xi1, #blocked0>) { %cst_0 = arith.constant dense<0.000000e+00> : tensor<256xf32, #blocked0> %1 = tt.load %a_ptr_init, %cst, %cst_0 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<256xf32, #blocked0> return@@ -72,7 +72,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #blocked0 = #triton_gpu.blocked<{sizePerThread = [4], threadsPerWarp = [32], warpsPerCTA = [8], order = [0]}> module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: masked_load_const_other_vec- func @masked_load_const_other_vec(%a_ptr_init : tensor<256x!tt.ptr<f32>, #blocked0>, %cst : tensor<256xi1, #blocked0>) {+ func.func @masked_load_const_other_vec(%a_ptr_init : tensor<256x!tt.ptr<f32>, #blocked0>, %cst : tensor<256xi1, #blocked0>) { %cst_0 = arith.constant dense<0.000000e+00> : tensor<256xf32, #blocked0> %1 = tt.load %a_ptr_init, %cst, %cst_0 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<256xf32, #blocked0> return@@ -84,7 +84,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #blocked0 = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [2], order = [0]}> module attributes {"triton_gpu.num-warps" = 2 : i32} { // CHECK-LABEL: global_load_store_no_vec- func @global_load_store_no_vec(%arg0: !tt.ptr<f32> {tt.divisibility = 4 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 4 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 4 : i32}, %arg3: i32) {+ func.func @global_load_store_no_vec(%arg0: !tt.ptr<f32> {tt.divisibility = 4 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 4 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 4 : i32}, %arg3: i32) { %c256_i32 = arith.constant 256 : i32 %0 = tt.get_program_id {axis = 0 : i32} : i32 %1 = arith.muli %0, %c256_i32 : i32@@ -128,7 +128,7 @@ module attributes {"triton_gpu.num-warps" = 2 : i32} { #blocked0 = #triton_gpu.blocked<{sizePerThread = [4], threadsPerWarp = [32], warpsPerCTA = [2], order = [0]}> module attributes {"triton_gpu.num-warps" = 2 : i32} { // CHECK-LABEL: global_load_store_vec4- func @global_load_store_vec4(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg3: i32) {+ func.func @global_load_store_vec4(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg3: i32) { %c256_i32 = arith.constant 256 : i32 %0 = tt.get_program_id {axis = 0 : i32} : i32 %1 = arith.muli %0, %c256_i32 : i32@@ -165,7 +165,7 @@ module attributes {"triton_gpu.num-warps" = 2 : i32} { #blocked = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [2], order = [0]}> // Note, the %n_elements doesn't have a "tt.divisibility" hint, so Triton assumes it's divisibility is 1, this should effect the mask's alignment and further restrict the load/store ops' vector width to be 1. module attributes {"triton_gpu.num-warps" = 2 : i32} {- func @vecadd_masked_vec1(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %n_elements: i32) {+ func.func @vecadd_masked_vec1(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %n_elements: i32) { %c64_i32 = arith.constant 64 : i32 %0 = tt.get_program_id {axis = 0 : i32} : i32 %1 = arith.muli %0, %c64_i32 : i32@@ -195,7 +195,7 @@ module attributes {"triton_gpu.num-warps" = 2 : i32} { #blocked0 = #triton_gpu.blocked<{sizePerThread = [8], threadsPerWarp = [32], warpsPerCTA = [1], order = [0]}> module attributes {"triton_gpu.num-warps" = 1 : i32} { // CHECK-LABEL: global_load_store_vec2- func @global_load_store_vec2(%arg0: !tt.ptr<f32> {tt.divisibility = 8 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 8 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 8 : i32}, %arg3: i32) {+ func.func @global_load_store_vec2(%arg0: !tt.ptr<f32> {tt.divisibility = 8 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 8 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 8 : i32}, %arg3: i32) { %c256_i32 = arith.constant 256 : i32 %0 = tt.get_program_id {axis = 0 : i32} : i32 %1 = arith.muli %0, %c256_i32 : i32@@ -240,7 +240,7 @@ module attributes {"triton_gpu.num-warps" = 1 : i32} { #blocked0 = #triton_gpu.blocked<{sizePerThread = [8], threadsPerWarp = [32], warpsPerCTA = [1], order = [0]}> module attributes {"triton_gpu.num-warps" = 1 : i32} { // CHECK-LABEL: global_load_store_vec8- func @global_load_store_vec8(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg3: i32) {+ func.func @global_load_store_vec8(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg3: i32) { %c256_i32 = arith.constant 256 : i32 %0 = tt.get_program_id {axis = 0 : i32} : i32 %1 = arith.muli %0, %c256_i32 : i32@@ -283,7 +283,7 @@ module attributes {"triton_gpu.num-warps" = 1 : i32} { #blocked2 = #triton_gpu.blocked<{sizePerThread = [1, 1], threadsPerWarp = [32, 1], warpsPerCTA = [4, 1], order = [0, 1]}> module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: basic_view_broadcast- func @basic_view_broadcast(%arg : tensor<256xf32,#blocked0>) {+ func.func @basic_view_broadcast(%arg : tensor<256xf32,#blocked0>) { // CHECK: llvm.mlir.undef // CHECK: %[[T0:.*]] = llvm.extractvalue // CHECK: %[[T1:.*]] = llvm.extractvalue@@ -307,7 +307,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #blocked0 = #triton_gpu.blocked<{sizePerThread = [2], threadsPerWarp = [32], warpsPerCTA = [1], order = [0]}> module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: basic_make_range- func @basic_make_range() {+ func.func @basic_make_range() { // CHECK: nvvm.read.ptx.sreg.tid.x // CHECK: llvm.mlir.undef // CHECK: llvm.insertvalue@@ -322,7 +322,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #blocked0 = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}> module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: basic_addf- func @basic_addf(%arg0 : tensor<256xf32,#blocked0>, %arg1 : tensor<256xf32,#blocked0>) {+ func.func @basic_addf(%arg0 : tensor<256xf32,#blocked0>, %arg1 : tensor<256xf32,#blocked0>) { // CHECK: llvm.fadd // CHECK: llvm.fadd %1 = arith.addf %arg0, %arg1 : tensor<256xf32,#blocked0>@@ -335,7 +335,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #blocked0 = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}> module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: basic_addi- func @basic_addi(%arg0 : tensor<256xi32,#blocked0>, %arg1 : tensor<256xi32,#blocked0>) {+ func.func @basic_addi(%arg0 : tensor<256xi32,#blocked0>, %arg1 : tensor<256xi32,#blocked0>) { // CHECK: llvm.add // CHECK: llvm.add %1 = arith.addi %arg0, %arg1 : tensor<256xi32,#blocked0>@@ -347,7 +347,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: basic_program_id- func @basic_program_id() {+ func.func @basic_program_id() { // CHECK: nvvm.read.ptx.sreg.ctaid.x : i32 %0 = tt.get_program_id {axis = 0 : i32} : i32 return@@ -359,7 +359,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #blocked0 = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}> module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: basic_addptr- func @basic_addptr(%arg0 : tensor<256x!tt.ptr<f32>,#blocked0>, %arg1 : tensor<256xi32,#blocked0>) {+ func.func @basic_addptr(%arg0 : tensor<256x!tt.ptr<f32>,#blocked0>, %arg1 : tensor<256xi32,#blocked0>) { // CHECK: llvm.getelementptr // CHECK: llvm.getelementptr %0 = tt.addptr %arg0, %arg1 : tensor<256x!tt.ptr<f32>, #blocked0>, tensor<256xi32, #blocked0>@@ -373,7 +373,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK: llvm.mlir.global external @global_smem // CHECK-LABEL: basic_alloc_tensor- func @basic_alloc_tensor() {+ func.func @basic_alloc_tensor() { // CHECK: llvm.mlir.addressof @global_smem // CHECK-NEXT: llvm.bitcast // CHECK-NEXT: llvm.mlir.constant@@ -390,7 +390,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK: llvm.mlir.global external @global_smem // CHECK-LABEL: basic_extract_slice- func @basic_extract_slice() {+ func.func @basic_extract_slice() { // CHECK: llvm.mlir.addressof @global_smem // CHECK: llvm.extractvalue // CHECK-NEXT: llvm.extractvalue@@ -423,7 +423,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: basic_async_wait- func @basic_async_wait() {+ func.func @basic_async_wait() { // CHECK: cp.async.wait_group 0x4 triton_gpu.async_wait {num = 4: i32} return@@ -442,7 +442,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #A = #triton_gpu.shared<{vec = 8, perPhase = 1, maxPhase = 4, order = [1, 0]}> module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: basic_insert_slice_async_fallback- func @basic_insert_slice_async_fallback(%arg0: !tt.ptr<f16> {tt.divisibility = 1 : i32}) {+ func.func @basic_insert_slice_async_fallback(%arg0: !tt.ptr<f16> {tt.divisibility = 1 : i32}) { %off0_ = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32, #slice2d1> %off1_ = tt.make_range {end = 32 : i32, start = 0 : i32} : tensor<64xi32, #slice3d0> %off0 = tt.expand_dims %off0_ {axis = 1 : i32} : (tensor<16xi32, #slice2d1>) -> tensor<16x1xi32, #block2>@@ -481,7 +481,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #A = #triton_gpu.shared<{vec = 8, perPhase = 1, maxPhase = 4, order = [1, 0]}> module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: basic_insert_slice_async_v4- func @basic_insert_slice_async_v4(%arg0: !tt.ptr<f32> {tt.divisibility = 32 : i32}) {+ func.func @basic_insert_slice_async_v4(%arg0: !tt.ptr<f32> {tt.divisibility = 32 : i32}) { %off0_ = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32, #slice2d1> %off1_ = tt.make_range {end = 32 : i32, start = 0 : i32} : tensor<64xi32, #slice3d0> %off0 = tt.expand_dims %off0_ {axis = 1 : i32} : (tensor<16xi32, #slice2d1>) -> tensor<16x1xi32, #block2>@@ -523,7 +523,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #A = #triton_gpu.shared<{vec = 1, perPhase = 1, maxPhase = 4, order = [1, 0]}> module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: basic_insert_slice_async_v1- func @basic_insert_slice_async_v1(%arg0: !tt.ptr<f32> {tt.divisibility = 4 : i32}) {+ func.func @basic_insert_slice_async_v1(%arg0: !tt.ptr<f32> {tt.divisibility = 4 : i32}) { %off0_ = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32, #slice2d1> %off1_ = tt.make_range {end = 32 : i32, start = 0 : i32} : tensor<32xi32, #slice3d0> %off0 = tt.expand_dims %off0_ {axis = 1 : i32} : (tensor<16xi32, #slice2d1>) -> tensor<16x1xi32, #block2>@@ -568,7 +568,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #A = #triton_gpu.shared<{vec = 1, perPhase = 1, maxPhase = 4, order = [1, 0]}> module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: basic_insert_slice_async_v1_multictas- func @basic_insert_slice_async_v1_multictas(%arg0: !tt.ptr<f32> {tt.divisibility = 4 : i32}) {+ func.func @basic_insert_slice_async_v1_multictas(%arg0: !tt.ptr<f32> {tt.divisibility = 4 : i32}) { %off0_ = tt.make_range {end = 32 : i32, start = 0 : i32} : tensor<32xi32, #slice2d1> %off1_ = tt.make_range {end = 32 : i32, start = 0 : i32} : tensor<32xi32, #slice3d0> %off0 = tt.expand_dims %off0_ {axis = 1 : i32} : (tensor<32xi32, #slice2d1>) -> tensor<32x1xi32, #block2>@@ -619,7 +619,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #blocked0 = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}> module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK: basic_splat- func @basic_splat(%ptr: !tt.ptr<f32>) {+ func.func @basic_splat(%ptr: !tt.ptr<f32>) { // CHECK: llvm.mlir.undef // CHECK: llvm.insertvalue // CHECK: llvm.insertvalue@@ -633,7 +633,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #blocked0 = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}> module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: basic_store- func @basic_store(%ptrs: tensor<256x!tt.ptr<f32>, #blocked0>, %vals: tensor<256xf32, #blocked0>, %mask: tensor<256xi1, #blocked0>) {+ func.func @basic_store(%ptrs: tensor<256x!tt.ptr<f32>, #blocked0>, %vals: tensor<256xf32, #blocked0>, %mask: tensor<256xi1, #blocked0>) { // CHECK: llvm.inline_asm // CHECK-SAME: st.global.b32 [ ${{.*}} + 0 ], { ${{.*}} }; // CHECK: llvm.inline_asm@@ -650,7 +650,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { module attributes {"triton_gpu.num-warps" = 1 : i32} { // CHECK: llvm.mlir.global external @global_smem() {addr_space = 3 : i32} : !llvm.array<0 x i8> // CHECK-LABEL: convert_layout_blocked_blocked- func @convert_layout_blocked_blocked(%arg0: tensor<16x16xf32, #blocked0>) {+ func.func @convert_layout_blocked_blocked(%arg0: tensor<16x16xf32, #blocked0>) { // CHECK: llvm.mlir.addressof @global_smem // CHECK: llvm.store // CHECK-SAME: !llvm.ptr<vector<1xf32>, 3>@@ -697,7 +697,7 @@ module attributes {"triton_gpu.num-warps" = 1 : i32} { module attributes {"triton_gpu.num-warps" = 1 : i32} { // CHECK: llvm.mlir.global external @global_smem() {addr_space = 3 : i32} : !llvm.array<0 x i8> // CHECK-LABEL: convert_layout_blocked_blocked_vec- func @convert_layout_blocked_blocked_vec(%arg0: tensor<16x16xf32, #blocked0>) {+ func.func @convert_layout_blocked_blocked_vec(%arg0: tensor<16x16xf32, #blocked0>) { // CHECK: llvm.mlir.addressof @global_smem // CHECK: llvm.store // CHECK-SAME: !llvm.ptr<vector<4xf32>, 3>@@ -720,7 +720,7 @@ module attributes {"triton_gpu.num-warps" = 1 : i32} { module attributes {"triton_gpu.num-warps" = 1 : i32} { // CHECK: llvm.mlir.global external @global_smem() {addr_space = 3 : i32} : !llvm.array<0 x i8> // CHECK-LABEL: convert_layout_blocked_blocked_multi_rep- func @convert_layout_blocked_blocked_multi_rep(%arg0: tensor<16x16xf32, #blocked0>) {+ func.func @convert_layout_blocked_blocked_multi_rep(%arg0: tensor<16x16xf32, #blocked0>) { // CHECK: llvm.mlir.addressof @global_smem // CHECK: llvm.store // CHECK-SAME: !llvm.ptr<vector<4xf32>, 3>@@ -751,7 +751,7 @@ module attributes {"triton_gpu.num-warps" = 1 : i32} { #dot_operand_b = #triton_gpu.dot_op<{opIdx=1, parent=#mma0}> module attributes {"triton_gpu.num-warps" = 1 : i32} { // CHECK-LABEL: convert_dot- func @convert_dot(%A: tensor<16x16xf16, #blocked0>, %B: tensor<16x16xf16, #blocked0>) {+ func.func @convert_dot(%A: tensor<16x16xf16, #blocked0>, %B: tensor<16x16xf16, #blocked0>) { %AA = triton_gpu.convert_layout %A : (tensor<16x16xf16, #blocked0>) -> tensor<16x16xf16, #shared0> %BB = triton_gpu.convert_layout %B : (tensor<16x16xf16, #blocked0>) -> tensor<16x16xf16, #shared0> // CHECK: llvm.inline_asm@@ -775,7 +775,7 @@ module attributes {"triton_gpu.num-warps" = 1 : i32} { // TODO: problems in MLIR's parser on slice layout // #blocked0 = #triton_gpu.blocked<{sizePerThread = [1, 4], threadsPerWarp = [8, 4], warpsPerCTA = [1, 1], order = [1, 0]}> // module attributes {"triton_gpu.num-warps" = 1 : i32} {-// func @make_range_sliced_layout() {+// func.func @make_range_sliced_layout() { // %0 = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32, #triton_gpu.slice<{dim = 0, parent = #blocked0}>> // return // }@@ -788,7 +788,7 @@ module attributes {"triton_gpu.num-warps" = 1 : i32} { module attributes {"triton_gpu.num-warps" = 1 : i32} { // CHECK: llvm.mlir.global external @global_smem() {addr_space = 3 : i32} : !llvm.array<0 x i8> // CHECK-LABEL: convert_layout_mmav2_block- func @convert_layout_mmav2_blocked(%arg0: tensor<32x16xf32, #mma>) {+ func.func @convert_layout_mmav2_blocked(%arg0: tensor<32x16xf32, #mma>) { // CHECK: llvm.store // CHECK-SAME: !llvm.ptr<vector<2xf32>, 3> // CHECK: llvm.store@@ -808,7 +808,7 @@ module attributes {"triton_gpu.num-warps" = 1 : i32} { module attributes {"triton_gpu.num-warps" = 1 : i32} { // CHECK: llvm.mlir.global external @global_smem() {addr_space = 3 : i32} : !llvm.array<0 x i8> // CHECK-LABEL: convert_layout_mmav1_block- func @convert_layout_mmav1_blocked(%arg0: tensor<32x64xf32, #mma>) {+ func.func @convert_layout_mmav1_blocked(%arg0: tensor<32x64xf32, #mma>) { // CHECK: llvm.store // CHECK-SAME: !llvm.ptr<vector<2xf32>, 3> // CHECK: llvm.store@@ -831,7 +831,7 @@ module attributes {"triton_gpu.num-warps" = 1 : i32} { module attributes {"triton_gpu.num-warps" = 1 : i32} { // CHECK: llvm.mlir.global external @global_smem() {addr_space = 3 : i32} : !llvm.array<0 x i8> // CHECK-LABEL: convert_layout_blocked_shared- func @convert_layout_blocked_shared(%arg0: tensor<128x32xf32, #blocked0>) {+ func.func @convert_layout_blocked_shared(%arg0: tensor<128x32xf32, #blocked0>) { // CHECK: llvm.store // CHECK-SAME: !llvm.ptr<vector<8xf32>, 3> // CHECK: llvm.store@@ -847,7 +847,7 @@ module attributes {"triton_gpu.num-warps" = 1 : i32} { #blocked1 = #triton_gpu.blocked<{sizePerThread = [1, 4], threadsPerWarp = [4, 8], warpsPerCTA = [1, 1], order = [1, 0]}> module attributes {"triton_gpu.num-warps" = 1 : i32} { // CHECK-LABEL: convert_blocked1d_to_slice0- func @convert_blocked1d_to_slice0(%src:tensor<32xi32, #blocked0>) {+ func.func @convert_blocked1d_to_slice0(%src:tensor<32xi32, #blocked0>) { // CHECK-COUNT-4: llvm.load {{.*}} : !llvm.ptr<vector<1xi32>, 3> %cvt = triton_gpu.convert_layout %src : (tensor<32xi32, #blocked0>) -> tensor<32xi32, #triton_gpu.slice<{dim = 0, parent = #blocked1}>> return@@ -860,7 +860,7 @@ module attributes {"triton_gpu.num-warps" = 1 : i32} { #blocked1 = #triton_gpu.blocked<{sizePerThread = [1, 4], threadsPerWarp = [4, 8], warpsPerCTA = [1, 1], order = [1, 0]}> module attributes {"triton_gpu.num-warps" = 1 : i32} { // CHECK-LABEL: convert_blocked1d_to_slice1- func @convert_blocked1d_to_slice1(%src:tensor<32xi32, #blocked0>) {+ func.func @convert_blocked1d_to_slice1(%src:tensor<32xi32, #blocked0>) { // CHECK-COUNT-32: llvm.load {{.*}} : !llvm.ptr<vector<1xi32>, 3> %cvt = triton_gpu.convert_layout %src : (tensor<32xi32, #blocked0>) -> tensor<32xi32, #triton_gpu.slice<{dim = 1, parent = #blocked1}>> return@@ -873,7 +873,7 @@ module attributes {"triton_gpu.num-warps" = 1 : i32} { #blocked1 = #triton_gpu.blocked<{sizePerThread = [4], threadsPerWarp = [32], warpsPerCTA = [1], order = [0]}> module attributes {"triton_gpu.num-warps" = 1 : i32} { // CHECK-LABEL: convert_blocked_to_blocked_ptr- func @convert_blocked_to_blocked_ptr(%src:tensor<32x!tt.ptr<f32>, #blocked0>) {+ func.func @convert_blocked_to_blocked_ptr(%src:tensor<32x!tt.ptr<f32>, #blocked0>) { // CHECK: llvm.ptrtoint // CHECK: llvm.store // CHECK: nvvm.barrier0@@ -892,7 +892,7 @@ module attributes {"triton_gpu.num-warps" = 1 : i32} { #dot_operand_a = #triton_gpu.dot_op<{opIdx=0, parent=#mma}> #dot_operand_b = #triton_gpu.dot_op<{opIdx=1, parent=#mma}> module attributes {"triton_gpu.num-warps" = 4 : i32} {- func @matmul_kernel_dot_operand_layout(%ptr:!tt.ptr<f32> {tt.divisibility = 16 : i32},+ func.func @matmul_kernel_dot_operand_layout(%ptr:!tt.ptr<f32> {tt.divisibility = 16 : i32}, %a:tensor<128x32xf16, #shared>, %b:tensor<32x256xf16, #shared>) { %cst = arith.constant dense<0.000000e+00> : tensor<128x256xf32, #mma> // CHECK: ldmatrix.sync.aligned.m8n8.x4.shared.b16@@ -918,7 +918,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #dot_operand_a = #triton_gpu.dot_op<{opIdx=0, parent=#mma, isMMAv1Row=true}> #dot_operand_b = #triton_gpu.dot_op<{opIdx=1, parent=#mma, isMMAv1Row=true}> module attributes {"triton_gpu.num-warps" = 4 : i32} {- func @matmul884_kernel_dot_operand_layout(%ptr:!tt.ptr<f32> {tt.divisibility = 16 : i32},+ func.func @matmul884_kernel_dot_operand_layout(%ptr:!tt.ptr<f32> {tt.divisibility = 16 : i32}, %a:tensor<32x64xf16, #shared0>, %b:tensor<64x64xf16, #shared1>) { %cst = arith.constant dense<0.000000e+00> : tensor<32x64xf32, #mma> // CHECK: ldmatrix.sync.aligned.m8n8.x4.shared.b16@@ -941,7 +941,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #dot_operand_a = #triton_gpu.dot_op<{opIdx=0, parent=#blocked}> #dot_operand_b = #triton_gpu.dot_op<{opIdx=1, parent=#blocked}> module attributes {"triton_gpu.num-warps" = 4 : i32} {- func @matmul_fmadot(%ptr:!tt.ptr<f32> {tt.divisibility = 16 : i32},+ func.func @matmul_fmadot(%ptr:!tt.ptr<f32> {tt.divisibility = 16 : i32}, %a:tensor<32x16xf32, #shared>, %b:tensor<16x32xf32, #shared>) { %cst = arith.constant dense<0.000000e+00> : tensor<32x32xf32, #blocked> // CHECK: llvm.intr.fmuladd@@ -965,7 +965,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #dot_operand_b = #triton_gpu.dot_op<{opIdx=1, parent=#mma}> module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: matmul_tf32dot- func @matmul_tf32dot(%ptr:!tt.ptr<f32> {tt.divisibility = 16 : i32},+ func.func @matmul_tf32dot(%ptr:!tt.ptr<f32> {tt.divisibility = 16 : i32}, %a:tensor<32x16xf32, #shared>, %b:tensor<16x32xf32, #shared>) { %cst = arith.constant dense<0.000000e+00> : tensor<32x32xf32, #mma> // CHECK: llvm.inline_asm@@ -1000,7 +1000,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #blocked0 = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}> module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: atomic_add_f32- func @atomic_add_f32(%arg0 : tensor<256x!tt.ptr<f32>, #blocked0>, %arg1 : tensor<256xi1, #blocked0>, %arg2 : tensor<256xf32, #blocked0>) {+ func.func @atomic_add_f32(%arg0 : tensor<256x!tt.ptr<f32>, #blocked0>, %arg1 : tensor<256xi1, #blocked0>, %arg2 : tensor<256xf32, #blocked0>) { // CHECK: llvm.inline_asm // CHECK-SAME: atom.global.gpu.add.f32 %0 = "tt.atomic_rmw" (%arg0, %arg2, %arg1) {atomic_rmw_op = 5 : i32} : (tensor<256x!tt.ptr<f32>, #blocked0>, tensor<256xf32, #blocked0>, tensor<256xi1, #blocked0>) -> tensor<256xf32, #blocked0>@@ -1012,7 +1012,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #blocked0 = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}> module attributes {"triton_gpu.num-warps" = 4 : i32} { -func @test_get_program_id(%a: tensor<32x!tt.ptr<i32>, #blocked0>) {+func.func @test_get_program_id(%a: tensor<32x!tt.ptr<i32>, #blocked0>) { %blockidx = tt.get_program_id {axis=0:i32} : i32 %blockidy = tt.get_program_id {axis=1:i32} : i32 %blockidz = tt.get_program_id {axis=2:i32} : i32@@ -1032,7 +1032,7 @@ func @test_get_program_id(%a: tensor<32x!tt.ptr<i32>, #blocked0>) { // ----- #blocked0 = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}> module attributes {"triton_gpu.num-warps" = 4 : i32} {- func @test_get_num_program(%a: tensor<32x!tt.ptr<i32>, #blocked0>) {+ func.func @test_get_num_program(%a: tensor<32x!tt.ptr<i32>, #blocked0>) { // CHECK: nvvm.read.ptx.sreg.nctaid.x // CHECK: nvvm.read.ptx.sreg.nctaid.y // CHECK: nvvm.read.ptx.sreg.nctaid.z@@ -1052,7 +1052,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #blocked0 = #triton_gpu.blocked<{sizePerThread = [2], threadsPerWarp = [32], warpsPerCTA = [1], order = [0]}> module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK-LABEL: test_index_cache- func @test_index_cache() {+ func.func @test_index_cache() { // CHECK: nvvm.read.ptx.sreg.tid.x %0 = tt.make_range {end = 256 : i32, start = 0 : i32} : tensor<256xi32, #blocked0> // CHECK-NOT: nvvm.read.ptx.sreg.tid.x@@ -1066,7 +1066,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { #shared0 = #triton_gpu.shared<{vec = 8, perPhase = 2, maxPhase = 4, order = [1, 0]}> module attributes {"triton_gpu.num-warps" = 1 : i32} { // CHECK-LABEL: test_base_index_cache- func @test_base_index_cache(%arg0: tensor<128x32xf32, #blocked0>) {+ func.func @test_base_index_cache(%arg0: tensor<128x32xf32, #blocked0>) { // CHECK: nvvm.read.ptx.sreg.tid.x %0 = triton_gpu.convert_layout %arg0 : (tensor<128x32xf32, #blocked0>) -> tensor<128x32xf32, #shared0> // CHECK-NOT: nvvm.read.ptx.sreg.tid.x@@ -1080,7 +1080,7 @@ module attributes {"triton_gpu.num-warps" = 1 : i32} { #shared0 = #triton_gpu.shared<{vec = 8, perPhase = 2, maxPhase = 4, order = [1, 0]}> module attributes {"triton_gpu.num-warps" = 1 : i32} { // CHECK-LABEL: test_index_cache_different_block- func @test_index_cache_different_block(%arg0: tensor<128x32xf32, #blocked0>, %arg1: i1) {+ func.func @test_index_cache_different_block(%arg0: tensor<128x32xf32, #blocked0>, %arg1: i1) { // CHECK: nvvm.read.ptx.sreg.tid.x %0 = triton_gpu.convert_layout %arg0 : (tensor<128x32xf32, #blocked0>) -> tensor<128x32xf32, #shared0> scf.if %arg1 {diff --git a/test/Target/tritongpu_to_llvmir.mlir b/test/Target/tritongpu_to_llvmir.mlirindex cafff3ca60..114d3a9eb2 100644--- a/test/Target/tritongpu_to_llvmir.mlir+++ b/test/Target/tritongpu_to_llvmir.mlir@@ -4,11 +4,11 @@ // CHECK-LABEL: ; ModuleID = 'LLVMDialectModule' // CHECK: define void @test_empty_kernel // CHECK: !nvvm.annotations-// CHECK: !{void (i32, half addrspace(1)*)* @test_empty_kernel, !"maxntidx", i32 128}+// CHECK: !{ptr @test_empty_kernel, !"maxntidx", i32 128} module attributes {"triton_gpu.num-warps" = 4 : i32} { -func @test_empty_kernel(%lb : index, %A : !tt.ptr<f16>) {+func.func @test_empty_kernel(%lb : index, %A : !tt.ptr<f16>) { return }diff --git a/test/Target/tritongpu_to_ptx.mlir b/test/Target/tritongpu_to_ptx.mlirindex 404e970a29..12742ad9e2 100644--- a/test/Target/tritongpu_to_ptx.mlir+++ b/test/Target/tritongpu_to_ptx.mlir@@ -6,7 +6,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { -func @test_empty_kernel(%lb : index, %A : !tt.ptr<f16>) {+func.func @test_empty_kernel(%lb : index, %A : !tt.ptr<f16>) { return }diff --git a/test/Triton/combine.mlir b/test/Triton/combine.mlirindex 050a3f7565..5ef6790e69 100644--- a/test/Triton/combine.mlir+++ b/test/Triton/combine.mlir@@ -2,10 +2,10 @@ // RUN: triton-opt %s -split-input-file -canonicalize -triton-combine | FileCheck %s // CHECK-LABEL: @test_combine_dot_add_pattern-func @test_combine_dot_add_pattern() -> (tensor<128x128xf32>, tensor<128x128xf32>) {- // CHECK: %[[d:.*]] = arith.constant dense<3.000000e+00> : tensor<128x128xf32>- // CHECK: %[[b:.*]] = arith.constant dense<2.000000e+00> : tensor<128x128xf32>- // CHECK: %[[a:.*]] = arith.constant dense<1.000000e+00> : tensor<128x128xf32>+func.func @test_combine_dot_add_pattern() -> (tensor<128x128xf32>, tensor<128x128xf32>) {+ // CHECK-DAG: %[[d:.*]] = arith.constant dense<3.000000e+00> : tensor<128x128xf32>+ // CHECK-DAG: %[[b:.*]] = arith.constant dense<2.000000e+00> : tensor<128x128xf32>+ // CHECK-DAG: %[[a:.*]] = arith.constant dense<1.000000e+00> : tensor<128x128xf32> %a = arith.constant dense<1.0> : tensor<128x128xf32> %b = arith.constant dense<2.0> : tensor<128x128xf32> %zero = arith.constant dense<0.0> : tensor<128x128xf32>@@ -24,7 +24,7 @@ func @test_combine_dot_add_pattern() -> (tensor<128x128xf32>, tensor<128x128xf32 // COM: CHECK-LABEL: @test_combine_addptr_pattern-func @test_combine_addptr_pattern(%base: !tt.ptr<f32>) -> tensor<8x!tt.ptr<f32>> {+func.func @test_combine_addptr_pattern(%base: !tt.ptr<f32>) -> tensor<8x!tt.ptr<f32>> { %off0 = arith.constant 10 : i32 %off1 = arith.constant 15 : i32 @@ -47,46 +47,46 @@ func @test_combine_addptr_pattern(%base: !tt.ptr<f32>) -> tensor<8x!tt.ptr<f32>> // CHECK-LABEL: @test_combine_select_masked_load_pattern-func @test_combine_select_masked_load_pattern(%ptr: tensor<8x!tt.ptr<f32>>, %cond: i1) -> (tensor<8xf32>, tensor<8xf32>) {+func.func @test_combine_select_masked_load_pattern(%ptr: tensor<8x!tt.ptr<f32>>, %cond: i1) -> (tensor<8xf32>, tensor<8xf32>) { %mask = tt.broadcast %cond : (i1) -> tensor<8xi1> %false_val = arith.constant dense<0.0> : tensor<8xf32> // CHECK: %[[res1:.*]] = tt.load %{{.*}}, %{{.*}}, %{{.*}} {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<8xf32> %x = tt.load %ptr, %mask, %false_val {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<8xf32>- %0 = select %cond, %x, %false_val : tensor<8xf32>+ %0 = arith.select %cond, %x, %false_val : tensor<8xf32> // CHECK: %[[res2:.*]] = tt.load %{{.*}}, %{{.*}}, %{{.*}} {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<8xf32> %y = tt.load %ptr, %mask, %false_val {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<8xf32>- %1 = select %cond, %y, %false_val : tensor<8xf32>+ %1 = arith.select %cond, %y, %false_val : tensor<8xf32> // CHECK: return %[[res1]], %[[res2]] : tensor<8xf32>, tensor<8xf32> return %0, %1 : tensor<8xf32>, tensor<8xf32> } // CHECK-LABEL: @test_combine_select_masked_load_fail_pattern-func @test_combine_select_masked_load_fail_pattern(%ptr: tensor<8x!tt.ptr<f32>>, %dummy_load: tensor<8xf32>, %dummy_broadcast: tensor<8xi1>, %cond0: i1, %cond1: i1) -> (tensor<8xf32>, tensor<8xf32>, tensor<8xf32>) {+func.func @test_combine_select_masked_load_fail_pattern(%ptr: tensor<8x!tt.ptr<f32>>, %dummy_load: tensor<8xf32>, %dummy_broadcast: tensor<8xi1>, %cond0: i1, %cond1: i1) -> (tensor<8xf32>, tensor<8xf32>, tensor<8xf32>) { %false_val = arith.constant dense<0.0> : tensor<8xf32> // Case 1: value at the "load" position is not an "op". Select should not be canonicalized.- // CHECK: %{{.*}} = select %{{.*}}, %{{.*}}, %{{.*}} : tensor<8xf32>- %0 = select %cond0, %dummy_load, %false_val : tensor<8xf32>+ // CHECK: %{{.*}} = arith.select %{{.*}}, %{{.*}}, %{{.*}} : tensor<8xf32>+ %0 = arith.select %cond0, %dummy_load, %false_val : tensor<8xf32> // Case 2: value at the "broadcast" position is not an "op". Select should not be canonicalized. %real_load0 = tt.load %ptr, %dummy_broadcast, %false_val {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<8xf32>- // CHECK: %{{.*}} = select %{{.*}}, %{{.*}}, %{{.*}} : tensor<8xf32>- %1 = select %cond0, %real_load0, %false_val : tensor<8xf32>+ // CHECK: %{{.*}} = arith.select %{{.*}}, %{{.*}}, %{{.*}} : tensor<8xf32>+ %1 = arith.select %cond0, %real_load0, %false_val : tensor<8xf32> // Case 3: condition of "broadcast" is not the same as the condition of "select". Select should not be canonicalized. %cond0_ = tt.broadcast %cond0 : (i1) -> tensor<8xi1> %real_load1 = tt.load %ptr, %cond0_, %false_val {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<8xf32>- // CHECK: %{{.*}} = select %{{.*}}, %{{.*}}, %{{.*}} : tensor<8xf32>- %2 = select %cond1, %real_load1, %false_val : tensor<8xf32>+ // CHECK: %{{.*}} = arith.select %{{.*}}, %{{.*}}, %{{.*}} : tensor<8xf32>+ %2 = arith.select %cond1, %real_load1, %false_val : tensor<8xf32> return %0, %1, %2 : tensor<8xf32>, tensor<8xf32>, tensor<8xf32> } // CHECK-LABEL: @test_combine_broadcast_constant_pattern-func @test_combine_broadcast_constant_pattern(%cst : f32) -> tensor<8x2xf32> {+func.func @test_combine_broadcast_constant_pattern(%cst : f32) -> tensor<8x2xf32> { // CHECK: %[[cst:.*]] = arith.constant dense<1.000000e+00> : tensor<8x2xf32> %const = arith.constant dense<1.0> : tensor<8xf32> %bst_out = tt.broadcast %const : (tensor<8xf32>) -> tensor<8x2xf32>@@ -96,7 +96,7 @@ func @test_combine_broadcast_constant_pattern(%cst : f32) -> tensor<8x2xf32> { } // CHECK-LABEL: @test_canonicalize_masked_load_pattern-func @test_canonicalize_masked_load_pattern(%ptr: tensor<8x!tt.ptr<f32>>) -> (tensor<8xf32>, tensor<8xf32>, tensor<8xf32>) {+func.func @test_canonicalize_masked_load_pattern(%ptr: tensor<8x!tt.ptr<f32>>) -> (tensor<8xf32>, tensor<8xf32>, tensor<8xf32>) { %true_mask = arith.constant dense<true> : tensor<8xi1> %false_mask = arith.constant dense<false> : tensor<8xi1> %other_val = arith.constant dense<0.0> : tensor<8xf32>@@ -117,7 +117,7 @@ func @test_canonicalize_masked_load_pattern(%ptr: tensor<8x!tt.ptr<f32>>) -> (te } // CHECK-LABEL: @test_canonicalize_masked_load_fail_pattern-func @test_canonicalize_masked_load_fail_pattern(%ptr: tensor<8x!tt.ptr<f32>>, %mask: tensor<8xi1>) -> (tensor<8xf32>, tensor<8xf32>) {+func.func @test_canonicalize_masked_load_fail_pattern(%ptr: tensor<8x!tt.ptr<f32>>, %mask: tensor<8xi1>) -> (tensor<8xf32>, tensor<8xf32>) { %other_val = arith.constant dense<0.0> : tensor<8xf32> // Case: value at the "mask" position is not an "op". Load should not be canonicalized.@@ -130,7 +130,7 @@ func @test_canonicalize_masked_load_fail_pattern(%ptr: tensor<8x!tt.ptr<f32>>, % } // CHECK-LABEL: @test_canonicalize_masked_store_pattern-func @test_canonicalize_masked_store_pattern(%ptr: tensor<8x!tt.ptr<f32>>, %val: tensor<8xf32>) {+func.func @test_canonicalize_masked_store_pattern(%ptr: tensor<8x!tt.ptr<f32>>, %val: tensor<8xf32>) { %true_mask = arith.constant dense<true> : tensor<8xi1> %false_mask = arith.constant dense<false> : tensor<8xi1> @@ -144,7 +144,7 @@ func @test_canonicalize_masked_store_pattern(%ptr: tensor<8x!tt.ptr<f32>>, %val: } // CHECK-LABEL: @test_canonicalize_masked_store_fail_pattern-func @test_canonicalize_masked_store_fail_pattern(%ptr: tensor<8x!tt.ptr<f32>>, %val: tensor<8xf32>, %mask: tensor<8xi1>) {+func.func @test_canonicalize_masked_store_fail_pattern(%ptr: tensor<8x!tt.ptr<f32>>, %val: tensor<8xf32>, %mask: tensor<8xi1>) { // Case: value at the "mask" position is not an "op". Store should not be canonicalized. // CHECK: tt.store %{{.*}}, %{{.*}}, %{{.*}} : tensor<8xf32> tt.store %ptr, %val, %mask : tensor<8xf32>diff --git a/test/Triton/vecadd.mlir b/test/Triton/vecadd.mlirindex 0b69ef3054..f5019b1cdd 100644--- a/test/Triton/vecadd.mlir+++ b/test/Triton/vecadd.mlir@@ -1,7 +1,7 @@ // RUN: triton-opt %s -verify-diagnostics module {- func @add_kernel__Pfp32_Pfp32_Pfp32_i32_i32_i32__(%arg0: !tt.ptr<f32>, %arg1: !tt.ptr<f32>, %arg2: !tt.ptr<f32>, %arg3: i32, %arg4: i32, %arg5: i32) {+ func.func @add_kernel__Pfp32_Pfp32_Pfp32_i32_i32_i32__(%arg0: !tt.ptr<f32>, %arg1: !tt.ptr<f32>, %arg2: !tt.ptr<f32>, %arg3: i32, %arg4: i32, %arg5: i32) { %0 = tt.get_program_id {axis = 0 : i32} : i32 %c256_i32 = arith.constant 256 : i32 %1 = arith.muli %0, %c256_i32 : i32@@ -43,7 +43,7 @@ module { } } // module {-// func @add_kernel__Pfp32_Pfp32_Pfp32_i32_i32_i32__(%arg0: !tt.ptr<f32>, %arg1: !tt.ptr<f32>, %arg2: !tt.ptr<f32>, %arg3: i32, %arg4: i32, %arg5: i32) {+// func.func @add_kernel__Pfp32_Pfp32_Pfp32_i32_i32_i32__(%arg0: !tt.ptr<f32>, %arg1: !tt.ptr<f32>, %arg2: !tt.ptr<f32>, %arg3: i32, %arg4: i32, %arg5: i32) { // %c64 = arith.constant 64 : index // %c32 = arith.constant 32 : index // %c0 = arith.constant 0 : indexdiff --git a/test/TritonGPU/coalesce.mlir b/test/TritonGPU/coalesce.mlirindex 60e359f527..51cccccfbd 100644--- a/test/TritonGPU/coalesce.mlir+++ b/test/TritonGPU/coalesce.mlir@@ -19,7 +19,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { // CHECK: [[store_val:%.*]] = triton_gpu.convert_layout {{.*}} -> tensor<64x64xf32, [[col_layout]]> // CHECK: [[store_mask:%.*]] = triton_gpu.convert_layout {{.*}} -> tensor<64x64xi1, [[col_layout]]> // CHECK: tt.store [[store_ptr]], [[store_val]], [[store_mask]]-func @transpose(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32},+func.func @transpose(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: i32 {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}) {diff --git a/test/TritonGPU/combine.mlir b/test/TritonGPU/combine.mlirindex 2c009ffa48..7e9cb9d504 100644--- a/test/TritonGPU/combine.mlir+++ b/test/TritonGPU/combine.mlir@@ -9,7 +9,7 @@ // CHECK: [[col_layout:#.*]] = #triton_gpu.blocked<{sizePerThread = [4, 1], threadsPerWarp = [16, 2], warpsPerCTA = [4, 1], order = [0, 1]}> // CHECK: [[col_layout_novec:#.*]] = #triton_gpu.blocked<{sizePerThread = [1, 1], threadsPerWarp = [32, 1], warpsPerCTA = [4, 1], order = [0, 1]}> // CHECK-LABEL: cst-func @cst() -> tensor<1024xi32, #layout1> {+func.func @cst() -> tensor<1024xi32, #layout1> { %cst = arith.constant dense<0> : tensor<1024xi32, #layout0> %1 = triton_gpu.convert_layout %cst : (tensor<1024xi32, #layout0>) -> tensor<1024xi32, #layout1> // CHECK-NOT: triton_gpu.convert_layout@@ -18,7 +18,7 @@ func @cst() -> tensor<1024xi32, #layout1> { } // CHECK-LABEL: range-func @range() -> tensor<1024xi32, #layout1> {+func.func @range() -> tensor<1024xi32, #layout1> { %0 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32, #layout0> %1 = triton_gpu.convert_layout %0 : (tensor<1024xi32, #layout0>) -> tensor<1024xi32, #layout1> // CHECK-NOT: triton_gpu.convert_layout@@ -27,7 +27,7 @@ func @range() -> tensor<1024xi32, #layout1> { } // CHECK-LABEL: splat-func @splat(%arg0: i32) -> tensor<1024xi32, #layout1> {+func.func @splat(%arg0: i32) -> tensor<1024xi32, #layout1> { %0 = tt.splat %arg0 : (i32) -> tensor<1024xi32, #layout0> %1 = triton_gpu.convert_layout %0 : (tensor<1024xi32, #layout0>) -> tensor<1024xi32, #layout1> // CHECK-NOT: triton_gpu.convert_layout@@ -36,7 +36,7 @@ func @splat(%arg0: i32) -> tensor<1024xi32, #layout1> { } // CHECK-LABEL: remat-func @remat(%arg0: i32) -> tensor<1024xi32, #layout1> {+func.func @remat(%arg0: i32) -> tensor<1024xi32, #layout1> { %0 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32, #layout0> %1 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32, #layout0> %2 = arith.muli %0, %1 : tensor<1024xi32, #layout0>@@ -56,7 +56,7 @@ func @remat(%arg0: i32) -> tensor<1024xi32, #layout1> { } // CHECK-LABEL: remat_load_store-func @remat_load_store(%arg: !tt.ptr<i32> {tt.divisibility = 16 : i32}) {+func.func @remat_load_store(%arg: !tt.ptr<i32> {tt.divisibility = 16 : i32}) { %0 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #layout0> %1 = tt.splat %arg : (!tt.ptr<i32>) -> tensor<64x!tt.ptr<i32>, #layout0> %2 = tt.addptr %1, %0 : tensor<64x!tt.ptr<i32>, #layout0>, tensor<64xi32, #layout0>@@ -70,7 +70,7 @@ func @remat_load_store(%arg: !tt.ptr<i32> {tt.divisibility = 16 : i32}) { // Don't rematerialize vectorized loads // CHECK-LABEL: remat_expensive-func @remat_expensive(%arg: !tt.ptr<i32> {tt.divisibility = 16 : i32}) {+func.func @remat_expensive(%arg: !tt.ptr<i32> {tt.divisibility = 16 : i32}) { %0 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #layout1> %1 = tt.splat %arg : (!tt.ptr<i32>) -> tensor<64x!tt.ptr<i32>, #layout1> %2 = tt.addptr %1, %0 : tensor<64x!tt.ptr<i32>, #layout1>, tensor<64xi32, #layout1>@@ -85,7 +85,7 @@ func @remat_expensive(%arg: !tt.ptr<i32> {tt.divisibility = 16 : i32}) { // Don't rematerialize loads when original and target layouts are different // CHECK-LABEL: remat_multi_layout-func @remat_multi_layout(%arg: !tt.ptr<i32> {tt.divisibility = 16 : i32}) {+func.func @remat_multi_layout(%arg: !tt.ptr<i32> {tt.divisibility = 16 : i32}) { %0 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #layout0> %1 = tt.splat %arg : (!tt.ptr<i32>) -> tensor<64x!tt.ptr<i32>, #layout0> %2 = tt.addptr %1, %0 : tensor<64x!tt.ptr<i32>, #layout0>, tensor<64xi32, #layout0>@@ -100,7 +100,7 @@ func @remat_multi_layout(%arg: !tt.ptr<i32> {tt.divisibility = 16 : i32}) { // Always rematerialize single value loads // CHECK-LABEL: remat_single_value-func @remat_single_value(%arg: !tt.ptr<i32> {tt.divisibility = 16 : i32}) {+func.func @remat_single_value(%arg: !tt.ptr<i32> {tt.divisibility = 16 : i32}) { %0 = tt.splat %arg : (!tt.ptr<i32>) -> tensor<1x!tt.ptr<i32>, #layout1> %1 = tt.load %0 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<1xi32, #layout1> // CHECK-NOT: triton_gpu.convert_layout@@ -111,7 +111,7 @@ func @remat_single_value(%arg: !tt.ptr<i32> {tt.divisibility = 16 : i32}) { } // CHECK-LABEL: if-func @if(%arg0: i32, %arg1: !tt.ptr<i32> {tt.divisibility = 16 : i32}) {+func.func @if(%arg0: i32, %arg1: !tt.ptr<i32> {tt.divisibility = 16 : i32}) { // CHECK-NOT: triton_gpu.convert_layout %c32_i32 = arith.constant dense<32> : tensor<1024xi32, #layout1> %0 = tt.get_program_id {axis = 0 : i32} : i32@@ -128,7 +128,7 @@ func @if(%arg0: i32, %arg1: !tt.ptr<i32> {tt.divisibility = 16 : i32}) { } // CHECK-LABEL: if_convert_else_not-func @if_convert_else_not(%arg0: i32, %arg1: !tt.ptr<i32> {tt.divisibility = 16 : i32}) {+func.func @if_convert_else_not(%arg0: i32, %arg1: !tt.ptr<i32> {tt.divisibility = 16 : i32}) { %c32_i32 = arith.constant dense<32> : tensor<1024xi32, #layout0> %0 = tt.get_program_id {axis = 0 : i32} : i32 %1 = tt.splat %0 : (i32) -> tensor<1024xi32, #layout0>@@ -149,7 +149,7 @@ func @if_convert_else_not(%arg0: i32, %arg1: !tt.ptr<i32> {tt.divisibility = 16 } // CHECK-LABEL: if_not_else_convert-func @if_not_else_convert(%arg0: i32, %arg1: !tt.ptr<i32> {tt.divisibility = 16 : i32}) {+func.func @if_not_else_convert(%arg0: i32, %arg1: !tt.ptr<i32> {tt.divisibility = 16 : i32}) { %c32_i32 = arith.constant dense<32> : tensor<1024xi32, #layout0> %0 = tt.get_program_id {axis = 0 : i32} : i32 %1 = tt.splat %0 : (i32) -> tensor<1024xi32, #layout0>@@ -170,7 +170,7 @@ func @if_not_else_convert(%arg0: i32, %arg1: !tt.ptr<i32> {tt.divisibility = 16 } // CHECK-LABEL: if_else_both_convert-func @if_else_both_convert(%arg0: i32, %arg1: !tt.ptr<i32> {tt.divisibility = 16 : i32}) {+func.func @if_else_both_convert(%arg0: i32, %arg1: !tt.ptr<i32> {tt.divisibility = 16 : i32}) { %c32_i32 = arith.constant dense<32> : tensor<1024xi32, #layout0> %0 = tt.get_program_id {axis = 0 : i32} : i32 %1 = tt.splat %0 : (i32) -> tensor<1024xi32, #layout0>@@ -200,7 +200,7 @@ func @if_else_both_convert(%arg0: i32, %arg1: !tt.ptr<i32> {tt.divisibility = 16 #blocked4 = #triton_gpu.blocked<{sizePerThread = [4, 1], threadsPerWarp = [16, 2], warpsPerCTA = [4, 1], order = [0, 1]}> // CHECK-LABEL: transpose-func @transpose(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: i32 {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}) {+func.func @transpose(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: i32 {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}) { // CHECK-NOT: triton_gpu.convert_layout // CHECK: [[loaded_val:%.*]] = tt.load {{.*}}, {{%cst.*}}, {{%cst.*}} {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<64x64xf32, [[row_layout]]> // CHECK: [[cvt_val:%.*]] = triton_gpu.convert_layout [[loaded_val]] : (tensor<64x64xf32, [[row_layout]]>) -> tensor<64x64xf32, [[col_layout]]>@@ -241,7 +241,7 @@ func @transpose(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: i32 {tt } // CHECK-LABEL: loop-func @loop(%arg0: !tt.ptr<f32>, %arg1: i32, %arg2: !tt.ptr<f32>, %arg3: i32, %arg4: i32) {+func.func @loop(%arg0: !tt.ptr<f32>, %arg1: i32, %arg2: !tt.ptr<f32>, %arg3: i32, %arg4: i32) { // CHECK-NOT: triton_gpu.convert_layout // CHECK: [[loop_ret:%.*]]:2 = scf.for {{.*}} -> (tensor<64x64xf32, [[row_layout]]>, tensor<64x64x!tt.ptr<f32>, [[row_layout]]>) // CHECK-NEXT: {{.*}} = tt.load {{.*}} : tensor<64x64xf32, [[row_layout]]>@@ -295,7 +295,7 @@ func @loop(%arg0: !tt.ptr<f32>, %arg1: i32, %arg2: !tt.ptr<f32>, %arg3: i32, %ar } // CHECK-LABEL: vecadd-func @vecadd(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg3: i32) {+func.func @vecadd(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg3: i32) { // CHECK-NOT: triton_gpu.convert_layout %c256_i32 = arith.constant 256 : i32 %0 = tt.get_program_id {axis = 0 : i32} : i32@@ -327,7 +327,7 @@ func @vecadd(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f3 // Select has args with different element types // CHECK-LABEL: select-func @select(%arg0: !tt.ptr<f64> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f64> {tt.divisibility = 16 : i32}, %arg2: i32 {tt.divisibility = 16 : i32}) {+func.func @select(%arg0: !tt.ptr<f64> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f64> {tt.divisibility = 16 : i32}, %arg2: i32 {tt.divisibility = 16 : i32}) { // CHECK-NOT: triton_gpu.convert_layout %cst = arith.constant dense<30000> : tensor<1x1xi32, #blocked2> %cst_0 = arith.constant dense<30000> : tensor<1x512xi32, #blocked2>@@ -378,7 +378,7 @@ func @select(%arg0: !tt.ptr<f64> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f6 // Make sure the following IR doesn't hang the compiler. // CHECK-LABEL: long_func-func public @long_func(%arg0: !tt.ptr<i64> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<i64> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<i64> {tt.divisibility = 16 : i32}, %arg3: !tt.ptr<i64> {tt.divisibility = 16 : i32}, %arg4: !tt.ptr<i64> {tt.divisibility = 16 : i32}, %arg5: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg6: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg7: !tt.ptr<i64> {tt.divisibility = 16 : i32}, %arg8: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg9: !tt.ptr<f64> {tt.divisibility = 16 : i32}, %arg10: !tt.ptr<f64> {tt.divisibility = 16 : i32}, %arg11: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg12: !tt.ptr<i32> {tt.divisibility = 16 : i32}, %arg13: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg14: !tt.ptr<f64> {tt.divisibility = 16 : i32}, %arg15: !tt.ptr<f64> {tt.divisibility = 16 : i32}, %arg16: i32 {tt.divisibility = 16 : i32}) {+func.func public @long_func(%arg0: !tt.ptr<i64> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<i64> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<i64> {tt.divisibility = 16 : i32}, %arg3: !tt.ptr<i64> {tt.divisibility = 16 : i32}, %arg4: !tt.ptr<i64> {tt.divisibility = 16 : i32}, %arg5: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg6: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg7: !tt.ptr<i64> {tt.divisibility = 16 : i32}, %arg8: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg9: !tt.ptr<f64> {tt.divisibility = 16 : i32}, %arg10: !tt.ptr<f64> {tt.divisibility = 16 : i32}, %arg11: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg12: !tt.ptr<i32> {tt.divisibility = 16 : i32}, %arg13: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg14: !tt.ptr<f64> {tt.divisibility = 16 : i32}, %arg15: !tt.ptr<f64> {tt.divisibility = 16 : i32}, %arg16: i32 {tt.divisibility = 16 : i32}) { %cst = arith.constant dense<1.000000e+00> : tensor<1024xf32, #blocked0> %cst_0 = arith.constant dense<5.000000e-04> : tensor<1024xf32, #blocked0> %cst_1 = arith.constant dense<0.999499976> : tensor<1024xf32, #blocked0>@@ -775,7 +775,7 @@ func public @long_func(%arg0: !tt.ptr<i64> {tt.divisibility = 16 : i32}, %arg1: // A mnist model from torch inductor. // Check if topological sort is working correct and there's no unnecessary convert // CHECK-LABEL: mnist-func public @mnist(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: i32 {tt.divisibility = 16 : i32}, %arg3: i32) {+func.func public @mnist(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: i32 {tt.divisibility = 16 : i32}, %arg3: i32) { // CHECK-NOT: triton_gpu.convert_layout %cst = arith.constant dense<10> : tensor<16x1xi32, #blocked2> %cst_0 = arith.constant dense<10> : tensor<1x16xi32, #blocked3>@@ -862,7 +862,7 @@ func public @mnist(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt. #blocked5 = #triton_gpu.blocked<{sizePerThread = [1, 4], threadsPerWarp = [2, 16], warpsPerCTA = [8, 1], order = [1, 0]}> // cmpf and cmpi have different operands and result types // CHECK-LABEL: cmp-func public @cmp(%arg0: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg3: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg4: i32 {tt.divisibility = 16 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}) {+func.func public @cmp(%arg0: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg3: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg4: i32 {tt.divisibility = 16 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}) { %c64 = arith.constant 64 : index %c2048 = arith.constant 2048 : index %c0 = arith.constant 0 : indexdiff --git a/test/TritonGPU/loop-pipeline.mlir b/test/TritonGPU/loop-pipeline.mlirindex 6ee3b15fbc..663f2da7b0 100644--- a/test/TritonGPU/loop-pipeline.mlir+++ b/test/TritonGPU/loop-pipeline.mlir@@ -10,7 +10,7 @@ #A = #triton_gpu.dot_op<{opIdx = 0, parent = #C}> #B = #triton_gpu.dot_op<{opIdx = 1, parent = #C}> -// CHECK: func @matmul_loop+// CHECK: func.func @matmul_loop // CHECK-DAG: %[[CONSTANT_0:.*]] = arith.constant 0 : i32 // CHECK-DAG: %[[CONSTANT_1:.*]] = arith.constant 1 : i32 // CHECK-DAG: %[[CONSTANT_2:.*]] = arith.constant 2 : i32@@ -46,8 +46,8 @@ // CHECK-DAG: %[[NEXT_PIPELINE_IDX:.*]] = arith.addi %[[PIPELINE_IDX]], %[[CONSTANT_1]] // CHECK-DAG: %[[NEXT_LOOP_IDX:.*]] = arith.addi %[[LOOP_IDX]], %[[CONSTANT_1]] // CHECK: scf.yield {{.*}}, {{.*}}, {{.*}}, %[[NEXT_A_BUFFER]], %[[NEXT_B_BUFFER]], %[[NEXT_A]], %[[NEXT_B]], {{.*}}, {{.*}}, {{.*}}, %[[NEXT_PIPELINE_IDX]], %[[NEXT_LOOP_IDX]]-func @matmul_loop(%lb : index, %ub : index, %step : index, - %A : !tt.ptr<f16> {tt.divisibility = 16 : i32}, +func.func @matmul_loop(%lb : index, %ub : index, %step : index,+ %A : !tt.ptr<f16> {tt.divisibility = 16 : i32}, %B : !tt.ptr<f16> {tt.divisibility = 16 : i32}) { // A ptrs %a_ptr_splat = tt.splat %A : (!tt.ptr<f16>) -> tensor<128x32x!tt.ptr<f16>, #AL>@@ -61,7 +61,7 @@ func @matmul_loop(%lb : index, %ub : index, %step : index, %b_tmp1 = tt.expand_dims %b_tmp0 {axis = 0 : i32} : (tensor<128xi32, #BLs0>) -> tensor<1x128xi32, #BL> %b_offs = tt.broadcast %b_tmp1 : (tensor<1x128xi32, #BL>) -> tensor<32x128xi32, #BL> %b_ptr_init = tt.addptr %b_ptr_splat, %b_offs : tensor<32x128x!tt.ptr<f16>, #BL>, tensor<32x128xi32, #BL>- + %a_mask = arith.constant dense<true> : tensor<128x32xi1, #AL> %a_other = arith.constant dense<0.00e+00> : tensor<128x32xf16, #AL>@@ -88,7 +88,7 @@ func @matmul_loop(%lb : index, %ub : index, %step : index, } -// CHECK: func @matmul_loop_nested+// CHECK: func.func @matmul_loop_nested // CHECK-DAG: %[[CONSTANT_0:.*]] = arith.constant 0 : i32 // CHECK-DAG: %[[CONSTANT_1:.*]] = arith.constant 1 : i32 // CHECK-DAG: %[[CONSTANT_2:.*]] = arith.constant 2 : i32@@ -118,8 +118,8 @@ func @matmul_loop(%lb : index, %ub : index, %step : index, // CHECK-DAG: %[[NEXT_PIPELINE_IDX:.*]] = arith.addi %[[PIPELINE_IDX]], %[[CONSTANT_1]] // CHECK-DAG: %[[NEXT_LOOP_IDX:.*]] = arith.addi %[[LOOP_IDX]], %[[CONSTANT_1]] // CHECK: scf.yield {{.*}}, {{.*}}, {{.*}}, %[[NEXT_A_BUFFER]], %[[NEXT_B_BUFFER]], %[[NEXT_A]], %[[NEXT_B]], {{.*}}, {{.*}}, {{.*}}, %[[NEXT_PIPELINE_IDX]], %[[NEXT_LOOP_IDX]]-func @matmul_loop_nested(%lb : index, %ub : index, %step : index, - %A : !tt.ptr<f16> {tt.divisibility = 16 : i32}, +func.func @matmul_loop_nested(%lb : index, %ub : index, %step : index,+ %A : !tt.ptr<f16> {tt.divisibility = 16 : i32}, %B : !tt.ptr<f16> {tt.divisibility = 16 : i32}) { scf.for %iv0 = %lb to %ub step %step { // A ptrs@@ -134,7 +134,7 @@ func @matmul_loop_nested(%lb : index, %ub : index, %step : index, %b_tmp1 = tt.expand_dims %b_tmp0 {axis = 0 : i32} : (tensor<128xi32, #BLs0>) -> tensor<1x128xi32, #BL> %b_offs = tt.broadcast %b_tmp1 : (tensor<1x128xi32, #BL>) -> tensor<32x128xi32, #BL> %b_ptr_init = tt.addptr %b_ptr_splat, %b_offs : tensor<32x128x!tt.ptr<f16>, #BL>, tensor<32x128xi32, #BL>- + %a_mask = arith.constant dense<true> : tensor<128x32xi1, #AL> %a_other = arith.constant dense<0.00e+00> : tensor<128x32xf16, #AL> %b_mask = arith.constant dense<true> : tensor<32x128xi1, #BL>@@ -161,7 +161,7 @@ func @matmul_loop_nested(%lb : index, %ub : index, %step : index, } -// CHECK: func @matmul_loop_single_pipeline+// CHECK: func.func @matmul_loop_single_pipeline // CHECK-DAG: %[[CONSTANT_0:.*]] = arith.constant 0 : i32 // CHECK-DAG: %[[CONSTANT_1:.*]] = arith.constant 1 : i32 // CHECK-DAG: %[[CONSTANT_2:.*]] = arith.constant 2 : i32@@ -183,8 +183,8 @@ func @matmul_loop_nested(%lb : index, %ub : index, %step : index, // CHECK-DAG: %[[NEXT_PIPELINE_IDX:.*]] = arith.addi %[[PIPELINE_IDX]], %[[CONSTANT_1]] // CHECK-DAG: %[[NEXT_LOOP_IDX:.*]] = arith.addi %[[LOOP_IDX]], %[[CONSTANT_1]] // CHECK: scf.yield {{.*}}, {{.*}}, %[[NEXT_B_BUFFER]], %[[NEXT_B]], {{.*}}, {{.*}}, %[[NEXT_PIPELINE_IDX]], %[[NEXT_LOOP_IDX]]-func @matmul_loop_single_pipeline(%lb : index, %ub : index, %step : index, - %A : !tt.ptr<f16> {tt.divisibility = 16 : i32}, +func.func @matmul_loop_single_pipeline(%lb : index, %ub : index, %step : index,+ %A : !tt.ptr<f16> {tt.divisibility = 16 : i32}, %B : !tt.ptr<f16> {tt.divisibility = 16 : i32}) { // A ptrs %a_ptr_splat = tt.splat %A : (!tt.ptr<f16>) -> tensor<128x32x!tt.ptr<f16>, #AL>diff --git a/test/TritonGPU/matmul.mlir b/test/TritonGPU/matmul.mlirindex 9bd5318e1e..01dc3f0ab1 100644--- a/test/TritonGPU/matmul.mlir+++ b/test/TritonGPU/matmul.mlir@@ -4,7 +4,7 @@ // CHECK: offset = 49152, size = 49152 // CHECK: size = 98304 module {-func @matmul_kernel__Pfp32_Pfp32_Pfp32_i32_i32_i32_i32_i32_i32_i32_i32_i32__12c64_13c64_14c64_15c8(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32 {tt.divisibility = 16 : i32}, %arg7: i32, %arg8: i32 {tt.divisibility = 16 : i32}, %arg9: i32, %arg10: i32 {tt.divisibility = 16 : i32}, %arg11: i32) {+func.func @matmul_kernel__Pfp32_Pfp32_Pfp32_i32_i32_i32_i32_i32_i32_i32_i32_i32__12c64_13c64_14c64_15c8(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32 {tt.divisibility = 16 : i32}, %arg7: i32, %arg8: i32 {tt.divisibility = 16 : i32}, %arg9: i32, %arg10: i32 {tt.divisibility = 16 : i32}, %arg11: i32) { %cst = arith.constant dense<true> : tensor<64x64xi1> %c64 = arith.constant 64 : index %c0 = arith.constant 0 : index@@ -22,7 +22,7 @@ func @matmul_kernel__Pfp32_Pfp32_Pfp32_i32_i32_i32_i32_i32_i32_i32_i32_i32__12c6 %7 = arith.muli %6, %c8_i32 : i32 %8 = arith.subi %2, %7 : i32 %9 = arith.cmpi slt, %8, %c8_i32 : i32- %10 = select %9, %8, %c8_i32 : i32+ %10 = arith.select %9, %8, %c8_i32 : i32 %11 = arith.remsi %0, %10 : i32 %12 = arith.addi %7, %11 : i32 %13 = arith.remsi %0, %5 : i32diff --git a/test/TritonGPU/prefetch.mlir b/test/TritonGPU/prefetch.mlirindex 52b4dddec1..b427547890 100644--- a/test/TritonGPU/prefetch.mlir+++ b/test/TritonGPU/prefetch.mlir@@ -11,7 +11,7 @@ #B_OP = #triton_gpu.dot_op<{opIdx = 1, parent = #C}> -// CHECK: func @matmul_loop+// CHECK: func.func @matmul_loop // CHECK-DAG: %[[A0_PREFETCH_SMEM:.*]] = tensor.extract_slice %[[A0:.*]][0, 0] [128, 16] // CHECK-DAG: %[[A0_PREFETCH:.*]] = triton_gpu.convert_layout %[[A0_PREFETCH_SMEM]] // CHECK-DAG: %[[B0_PREFETCH_SMEM:.*]] = tensor.extract_slice %[[B0:.*]][0, 0] [16, 128]@@ -28,7 +28,7 @@ // CHECK-DAG: %[[NEXT_B_PREFETCH_SMEM:.*]] = tensor.extract_slice {{.*}}[0, 0] [16, 128] // CHECK-DAG: %[[NEXT_B_PREFETCH:.*]] = triton_gpu.convert_layout %[[NEXT_B_PREFETCH_SMEM]] // CHECK: scf.yield {{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}, %[[NEXT_A_PREFETCH]], %[[NEXT_B_PREFETCH]]-func @matmul_loop(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) {+func.func @matmul_loop(%lb : index, %ub : index, %step : index, %A : !tt.ptr<f16>, %B : !tt.ptr<f16>) { %a_ptr_init = tt.broadcast %A : (!tt.ptr<f16>) -> tensor<128x32x!tt.ptr<f16>, #AL> %b_ptr_init = tt.broadcast %B : (!tt.ptr<f16>) -> tensor<32x128x!tt.ptr<f16>, #BL> diff --git a/test/TritonGPU/update-mma-for-volta.mlir b/test/TritonGPU/update-mma-for-volta.mlirindex d587fffcca..7571ec6185 100644--- a/test/TritonGPU/update-mma-for-volta.mlir+++ b/test/TritonGPU/update-mma-for-volta.mlir@@ -15,7 +15,7 @@ // CHECK: [[new_mma:#mma.*]] = #triton_gpu.mma<{versionMajor = 1, versionMinor = 3, warpsPerCTA = [4, 2]}> module attributes {"triton_gpu.num-warps" = 16 : i32} { // CHECK-LABEL: dot_mmav1- func @dot_mmav1(%A: tensor<64x64xf16, #blocked0>, %B: tensor<64x64xf16, #blocked0>) -> tensor<64x64xf32, #blocked0> {+ func.func @dot_mmav1(%A: tensor<64x64xf16, #blocked0>, %B: tensor<64x64xf16, #blocked0>) -> tensor<64x64xf32, #blocked0> { %C = arith.constant dense<0.000000e+00> : tensor<64x64xf32, #blocked0> %AA = triton_gpu.convert_layout %A : (tensor<64x64xf16, #blocked0>) -> tensor<64x64xf16, #dot_operand_a> %BB = triton_gpu.convert_layout %B : (tensor<64x64xf16, #blocked0>) -> tensor<64x64xf16, #dot_operand_b>@@ -50,7 +50,7 @@ module attributes {"triton_gpu.num-warps" = 16 : i32} { module attributes {"triton_gpu.num-warps" = 16 : i32} { // CHECK-LABEL: dot_mmav1- func @dot_mmav1(%A: tensor<64x64xf16, #blocked0>, %B: tensor<64x64xf16, #blocked0>) -> tensor<64x64xf32, #blocked0> {+ func.func @dot_mmav1(%A: tensor<64x64xf16, #blocked0>, %B: tensor<64x64xf16, #blocked0>) -> tensor<64x64xf32, #blocked0> { %C = arith.constant dense<0.000000e+00> : tensor<64x64xf32, #blocked0> %AA = triton_gpu.convert_layout %A : (tensor<64x64xf16, #blocked0>) -> tensor<64x64xf16, #dot_operand_a> %BB = triton_gpu.convert_layout %B : (tensor<64x64xf16, #blocked0>) -> tensor<64x64xf16, #dot_operand_b>diff --git a/test/lib/Analysis/TestAlias.cpp b/test/lib/Analysis/TestAlias.cppindex 88a4118fe9..3fd0cfd0d3 100644--- a/test/lib/Analysis/TestAlias.cpp+++ b/test/lib/Analysis/TestAlias.cpp@@ -9,10 +9,10 @@ using namespace mlir; namespace { struct TestAliasPass- : public PassWrapper<TestAliasPass, OperationPass<FuncOp>> {+ : public PassWrapper<TestAliasPass, OperationPass<func::FuncOp>> {++ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAliasPass); - // LLVM15+- // MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAliasPass); static void print(StringRef name, SmallVector<std::string, 4> &vals, raw_ostream &os) { if (vals.empty())@@ -39,23 +39,24 @@ struct TestAliasPass auto opName = SymbolTable::getSymbolName(operation).getValue().str(); os << opName << "\n"; - SharedMemoryAliasAnalysis analysis(&getContext());- analysis.run(operation);+ std::unique_ptr<DataFlowSolver> solver = createDataFlowSolver();+ SharedMemoryAliasAnalysis *analysis =+ solver->load<SharedMemoryAliasAnalysis>();+ if (failed(solver->initializeAndRun(operation)))+ return signalPassFailure(); AsmState state(operation->getParentOfType<ModuleOp>()); // Get operation ids of value's aliases auto getAllocOpNames = [&](Value value) {- LatticeElement<AliasInfo> *latticeElement =- analysis.lookupLatticeElement(value);+ dataflow::Lattice<AliasInfo> *latticeElement =+ analysis->getLatticeElement(value); SmallVector<std::string, 4> opNames;- if (latticeElement) {+ if (latticeElement && !latticeElement->isUninitialized()) { auto &info = latticeElement->getValue();- if (!info.getAllocs().empty()) {- for (auto &alias : info.getAllocs()) {- auto opName =- getValueOperandName(alias.getDefiningOp()->getResult(0), state);- opNames.push_back(std::move(opName));- }+ for (auto &alias : info.getAllocs()) {+ auto opName =+ getValueOperandName(alias.getDefiningOp()->getResult(0), state);+ opNames.push_back(std::move(opName)); } } // Ensure deterministic outputdiff --git a/test/lib/Analysis/TestAllocation.cpp b/test/lib/Analysis/TestAllocation.cppindex 84108c4d36..35e42242bd 100644--- a/test/lib/Analysis/TestAllocation.cpp+++ b/test/lib/Analysis/TestAllocation.cpp@@ -6,10 +6,9 @@ using namespace mlir; namespace { struct TestAllocationPass- : public PassWrapper<TestAllocationPass, OperationPass<FuncOp>> {+ : public PassWrapper<TestAllocationPass, OperationPass<func::FuncOp>> { - // LLVM15+- // MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAllocationPass);+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAllocationPass); StringRef getArgument() const final { return "test-print-allocation"; } StringRef getDescription() const final {diff --git a/test/lib/Analysis/TestAxisInfo.cpp b/test/lib/Analysis/TestAxisInfo.cppindex a5205bb0a0..22347c32f0 100644--- a/test/lib/Analysis/TestAxisInfo.cpp+++ b/test/lib/Analysis/TestAxisInfo.cpp@@ -1,25 +1,15 @@ #include "mlir/Pass/Pass.h" #include "triton/Analysis/AxisInfo.h"+#include "triton/Analysis/Utility.h" using namespace mlir; namespace { struct TestAxisInfoPass- : public PassWrapper<TestAxisInfoPass, OperationPass<FuncOp>> {+ : public PassWrapper<TestAxisInfoPass, OperationPass<func::FuncOp>> { - // LLVM15+- // MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAlignmentPass);-- void print(const std::string &name, raw_ostream &os, ArrayRef<int64_t> vals) {- os << name << ": [";- for (size_t d = 0; d < vals.size(); d++) {- if (d != 0)- os << ", ";- os << vals[d];- }- os << "]";- }+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAxisInfoPass); StringRef getArgument() const final { return "test-print-alignment"; } StringRef getDescription() const final {@@ -30,38 +20,19 @@ struct TestAxisInfoPass Operation *operation = getOperation(); auto &os = llvm::errs(); auto opName = SymbolTable::getSymbolName(operation).getValue().str();- os << opName << "\n";- AxisInfoAnalysis analysis(&getContext());- analysis.run(operation);+ os << "@" << opName << "\n";++ std::unique_ptr<DataFlowSolver> solver = createDataFlowSolver();+ AxisInfoAnalysis *analysis = solver->load<AxisInfoAnalysis>();+ if (failed(solver->initializeAndRun(operation)))+ return signalPassFailure(); operation->walk([&](Operation *op) { if (op->getNumResults() < 1) return; for (Value result : op->getResults()) {- // std::ostringstream oss;- // result.print(oss);- // os << " => ";- LatticeElement<AxisInfo> *latticeElement =- analysis.lookupLatticeElement(result);- if (!latticeElement) {- os << "None\n";- return;- }- AxisInfo &info = latticeElement->getValue();- print("Contiguity", os, info.getContiguity());- os << " ; ";- print("Divisibility", os, info.getDivisibility());- os << " ; ";- print("Constancy", os, info.getConstancy());- os << " ; ";- auto constantValue = info.getConstantValue();- os << "ConstantValue: [";- if (constantValue.has_value())- os << constantValue.value();- else- os << "None";- os << "] ( "; result.print(os);- os << " ) ";+ os << " => ";+ analysis->getLatticeElement(result)->getValue().print(os); os << "\n"; } });diff --git a/test/lib/Analysis/TestMembar.cpp b/test/lib/Analysis/TestMembar.cppindex df4279fe24..ab9b9f3fb7 100644--- a/test/lib/Analysis/TestMembar.cpp+++ b/test/lib/Analysis/TestMembar.cpp@@ -1,4 +1,4 @@-#include "mlir/Dialect/GPU/GPUDialect.h"+#include "mlir/Dialect/GPU/IR/GPUDialect.h" #include "mlir/IR/Dialect.h" #include "mlir/Pass/Pass.h" #include "triton/Analysis/Allocation.h"@@ -9,10 +9,9 @@ using namespace mlir; namespace { struct TestMembarPass- : public PassWrapper<TestMembarPass, OperationPass<FuncOp>> {+ : public PassWrapper<TestMembarPass, OperationPass<func::FuncOp>> { - // LLVM15+- // MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestMembarPass);+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestMembarPass); StringRef getArgument() const final { return "test-print-membar"; } StringRef getDescription() const final {